fix imports

server/text_generation_server/models/flash_causal_lm.py

@@ -18,10 +18,7 @@ from text_generation_server.models.types import (
     GeneratedText,
 )
 from text_generation_server.pb import generate_pb2
-from text_generation_server.utils import (
+from text_generation_server.utils import StoppingCriteria, HeterogeneousNextTokenChooser
-    StoppingCriteria,
-    HeterogeneousNextTokenChooser
-)
 
 tracer = trace.get_tracer(__name__)
 
@@ -228,7 +225,7 @@ class FlashCausalLMBatch(Batch):
 
             # Slice from past
             past_key_values.append(
-                self.past_key_values[:, self.cu_seqlens[idx]: self.cu_seqlens[idx + 1]]
+                self.past_key_values[:, self.cu_seqlens[idx] : self.cu_seqlens[idx + 1]]
             )
 
             all_input_ids.append(self.all_input_ids[idx])
@@ -630,8 +627,8 @@ class FlashCausalLM(Model):
                 # Copy batch.input_ids to prefill_token_indices
                 if len(batch) > 1:
                     prefill_tokens_indices[
-                    start_index: end_index - 1
+                        start_index : end_index - 1
-                    ] = batch.input_ids[start_index + 1: end_index]
+                    ] = batch.input_ids[start_index + 1 : end_index]
                 else:
                     # Set prefill_tokens_indices to the correct slice
                     prefill_tokens_indices = batch.input_ids
@@ -717,7 +714,7 @@ class FlashCausalLM(Model):
                 if stop:
                     # Decode generated tokens
                     output_text = self.decode(
-                        all_input_ids[-stopping_criteria.current_tokens:]
+                        all_input_ids[-stopping_criteria.current_tokens :]
                     )
                     generated_text = GeneratedText(
                         output_text,
@@ -732,7 +729,7 @@ class FlashCausalLM(Model):
                 if prefill:
                     # Remove generated token to only have prefill and add nan for first prompt token
                     request_prefill_logprobs = [float("nan")] + prefill_logprobs[
-                                                                start_index: end_index - 1
+                        start_index : end_index - 1
                     ]
                     prefill_token_ids = all_input_ids[:-1]
                     prefill_texts = self.tokenizer.batch_decode(

server/text_generation_server/utils/__init__.py

@@ -14,8 +14,9 @@ from text_generation_server.utils.tokens import (
     StoppingCriteria,
     StopSequenceCriteria,
     FinishReason,
+    Sampling,
+    Greedy,
 )
-from text_generation_server.utils.logits_process import Sampling, Greedy
 
 __all__ = [
     "convert_file",

server/text_generation_server/utils/logits_process.py

@@ -14,25 +14,6 @@ from transformers import (
 )
 
 
-class Sampling:
-    def __init__(self, seed: int, device: str = "cpu"):
-        self.generator = torch.Generator(device)
-        self.generator.manual_seed(seed)
-        self.seed = seed
-
-    def __call__(self, logits):
-        probs = torch.nn.functional.softmax(logits, -1)
-        # Avoid GPU<->CPU sync done by torch multinomial
-        # See: https://github.com/pytorch/pytorch/blob/925a3788ec5c06db62ca732a0e9425a26a00916f/aten/src/ATen/native/Distributions.cpp#L631-L637
-        q = torch.empty_like(probs).exponential_(1, generator=self.generator)
-        return probs.div_(q).argmax()
-
-
-class Greedy:
-    def __call__(self, logits):
-        return logits.argmax(dim=-1)
-
-
 class StaticWarper:
     def __init__(
         self,
@@ -329,46 +310,3 @@ class HeterogeneousTypicalLogitsWarper(LogitsWarper):
     def filter(self, indices):
         self.mass = self.mass[indices]
         return self
-
-
-class HeterogeneousSampling:
-    r"""
-    Mixed greedy and probabilistic sampling. Compute both and pick the right one for each sample.
-    """
-
-    def __init__(self, do_sample: List[bool], seeds: List[int], device: torch.device):
-        self.seeds = seeds
-
-        self.greedy_indices = []
-        self.sampling_mapping = {}
-        for i, (sample, seed) in enumerate(zip(do_sample, seeds)):
-            if sample:
-                self.sampling_mapping[i] = Sampling(seed, device)
-            else:
-                self.greedy_indices.append(i)
-
-        self.greedy = Greedy()
-
-    def __call__(self, logits):
-        out = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device)
-        if self.greedy_indices:
-            out[self.greedy_indices] = torch.argmax(logits[self.greedy_indices], -1)
-
-        for i, sampling in self.sampling_mapping.items():
-            out[i] = sampling(logits[i])
-        return out
-
-    def filter(self, indices):
-        new_greedy_indices = []
-        new_sampling_mapping = {}
-        for i, idx in enumerate(indices):
-            if idx in self.sampling_mapping:
-                new_sampling_mapping[i] = self.sampling_mapping[idx]
-            else:
-                new_greedy_indices.append(i)
-
-        self.greedy_indices = new_greedy_indices
-        self.sampling_mapping = new_sampling_mapping
-        return self
-
-

server/text_generation_server/utils/tokens.py

@@ -3,17 +3,22 @@ import torch
 
 from transformers import (
     RepetitionPenaltyLogitsProcessor,
-    PreTrainedTokenizerBase, LogitsProcessorList,
+    PreTrainedTokenizerBase,
+    LogitsProcessorList,
 )
 from typing import List, Tuple, Optional
 
 from text_generation_server.pb import generate_pb2
 from text_generation_server.pb.generate_pb2 import FinishReason
 from text_generation_server.utils.watermark import WatermarkLogitsProcessor
-from text_generation_server.utils import Sampling, Greedy
+from text_generation_server.utils.logits_process import (
-from text_generation_server.utils.logits_process import static_warper, HeterogeneousRepetitionPenaltyLogitsProcessor, \
+    static_warper,
-    HeterogeneousTemperatureLogitsWarper, HeterogeneousTopKLogitsWarper, HeterogeneousTopPLogitsWarper, \
+    HeterogeneousRepetitionPenaltyLogitsProcessor,
-    HeterogeneousTypicalLogitsWarper, HeterogeneousSampling
+    HeterogeneousTemperatureLogitsWarper,
+    HeterogeneousTopKLogitsWarper,
+    HeterogeneousTopPLogitsWarper,
+    HeterogeneousTypicalLogitsWarper,
+)
 
 
 class NextTokenChooser:
@@ -240,3 +245,63 @@ class HeterogeneousNextTokenChooser:
             device=device,
             dtype=dtype,
         )
+
+
+class Sampling:
+    def __init__(self, seed: int, device: str = "cpu"):
+        self.generator = torch.Generator(device)
+        self.generator.manual_seed(seed)
+        self.seed = seed
+
+    def __call__(self, logits):
+        probs = torch.nn.functional.softmax(logits, -1)
+        # Avoid GPU<->CPU sync done by torch multinomial
+        # See: https://github.com/pytorch/pytorch/blob/925a3788ec5c06db62ca732a0e9425a26a00916f/aten/src/ATen/native/Distributions.cpp#L631-L637
+        q = torch.empty_like(probs).exponential_(1, generator=self.generator)
+        return probs.div_(q).argmax()
+
+
+class Greedy:
+    def __call__(self, logits):
+        return logits.argmax(dim=-1)
+
+
+class HeterogeneousSampling:
+    r"""
+    Mixed greedy and probabilistic sampling. Compute both and pick the right one for each sample.
+    """
+
+    def __init__(self, do_sample: List[bool], seeds: List[int], device: torch.device):
+        self.seeds = seeds
+
+        self.greedy_indices = []
+        self.sampling_mapping = {}
+        for i, (sample, seed) in enumerate(zip(do_sample, seeds)):
+            if sample:
+                self.sampling_mapping[i] = Sampling(seed, device)
+            else:
+                self.greedy_indices.append(i)
+
+        self.greedy = Greedy()
+
+    def __call__(self, logits):
+        out = torch.empty(logits.shape[0], dtype=torch.int64, device=logits.device)
+        if self.greedy_indices:
+            out[self.greedy_indices] = torch.argmax(logits[self.greedy_indices], -1)
+
+        for i, sampling in self.sampling_mapping.items():
+            out[i] = sampling(logits[i])
+        return out
+
+    def filter(self, indices):
+        new_greedy_indices = []
+        new_sampling_mapping = {}
+        for i, idx in enumerate(indices):
+            if idx in self.sampling_mapping:
+                new_sampling_mapping[i] = self.sampling_mapping[idx]
+            else:
+                new_greedy_indices.append(i)
+
+        self.greedy_indices = new_greedy_indices
+        self.sampling_mapping = new_sampling_mapping
+        return self