Ralf12358
diff --git a/‎llama_cpp/llama.py
Copy file name to clipboardExpand all lines: llama_cpp/llama.py
+26Lines changed: 26 additions & 0 deletions b/‎llama_cpp/llama.py
Copy file name to clipboardExpand all lines: llama_cpp/llama.py
+26Lines changed: 26 additions & 0 deletions
@@ -290,7 +290,11 @@ def _sample(
         mirostat_tau: llama_cpp.c_float,
         mirostat_eta: llama_cpp.c_float,
         penalize_nl: bool = True,
+        logits_processors=None
     ):
+        if logits_processors is None:
+            logits_processors = []
+
         assert self.ctx is not None
         assert len(self.eval_logits) > 0
         n_vocab = int(llama_cpp.llama_n_vocab(self.ctx))
@@ -302,6 +306,9 @@ def _sample(
             else last_n_tokens_size
         )
         logits = self.eval_logits[-1]
+        for processor in logits_processors:
+            logits = processor(last_n_tokens_data, logits)
+
         nl_logit = logits[int(Llama.token_nl())]
         data = (llama_cpp.llama_token_data * n_vocab)(
             *[
@@ -420,6 +427,7 @@ def sample(
         mirostat_eta: float = 0.1,
         mirostat_tau: float = 5.0,
         penalize_nl: bool = True,
+        logits_processors=None
     ):
         """Sample a token from the model.
 
@@ -452,6 +460,7 @@ def sample(
             mirostat_tau=llama_cpp.c_float(mirostat_tau),
             mirostat_eta=llama_cpp.c_float(mirostat_eta),
             penalize_nl=penalize_nl,
+            logits_processors=logits_processors
         )
 
     def generate(
@@ -468,6 +477,7 @@ def generate(
         mirostat_mode: int = 0,
         mirostat_tau: float = 5.0,
         mirostat_eta: float = 0.1,
+        logits_processors=None
     ) -> Generator[int, Optional[Sequence[int]], None]:
         """Create a generator of tokens from a prompt.
 
@@ -525,6 +535,7 @@ def generate(
                 mirostat_mode=mirostat_mode,
                 mirostat_tau=mirostat_tau,
                 mirostat_eta=mirostat_eta,
+                logits_processors=logits_processors
             )
             tokens_or_none = yield token
             tokens = [token]
@@ -609,6 +620,8 @@ def _create_completion(
         mirostat_tau: float = 5.0,
         mirostat_eta: float = 0.1,
         model: Optional[str] = None,
+        logits_processors=None,
+        stopping_criterias=None
     ) -> Union[Iterator[Completion], Iterator[CompletionChunk]]:
         assert self.ctx is not None
         completion_id: str = f"cmpl-{str(uuid.uuid4())}"
@@ -672,13 +685,22 @@ def _create_completion(
             frequency_penalty=frequency_penalty,
             presence_penalty=presence_penalty,
             repeat_penalty=repeat_penalty,
+            logits_processors=logits_processors
         ):
             if token == Llama.token_eos():
                 text = self.detokenize(completion_tokens)
                 finish_reason = "stop"
                 break
 
             completion_tokens.append(token)
+            for stopping_crit in stopping_criterias:
+                if stopping_crit(completion_tokens, None):
+                    text = self.detokenize(completion_tokens)
+                    finish_reason = "stop"
+                    break
+
+            if finish_reason == "stop":
+                break
 
             all_text = self.detokenize(completion_tokens)
 
@@ -978,6 +1000,8 @@ def create_completion(
         mirostat_tau: float = 5.0,
         mirostat_eta: float = 0.1,
         model: Optional[str] = None,
+        logits_processors=None,
+        stopping_criterias=None
     ) -> Union[Completion, Iterator[CompletionChunk]]:
         """Generate text from a prompt.
 
@@ -1020,6 +1044,8 @@ def create_completion(
             mirostat_tau=mirostat_tau,
             mirostat_eta=mirostat_eta,
             model=model,
+            logits_processors=logits_processors,
+            stopping_criterias=stopping_criterias
         )
         if stream:
             chunks: Iterator[CompletionChunk] = completion_or_chunks