bugfix: truncate completion max_tokens to fit context length by default

abetlen · abetlen · commit a86bfdf0a50f · 2023-07-09T18:13:29.000-04:00
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -824,19 +824,15 @@ def _create_completion(
         if self.verbose:
             llama_cpp.llama_reset_timings(self.ctx)
 
-        if max_tokens <= 0:
-            # Unlimited, depending on n_ctx.
-            if len(prompt_tokens) >= int(llama_cpp.llama_n_ctx(self.ctx)):
-                raise ValueError(
-                    f"Requested tokens exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}"
-                )
-            else:
-                max_tokens = int(llama_cpp.llama_n_ctx(self.ctx)) - len(prompt_tokens)
-        elif len(prompt_tokens) + max_tokens > int(llama_cpp.llama_n_ctx(self.ctx)):
+        if len(prompt_tokens) >= llama_cpp.llama_n_ctx(self.ctx):
             raise ValueError(
-                f"Requested tokens ({len(prompt_tokens)}) exceed context window of {self._n_ctx}"
+                f"Requested tokens exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}"
             )
 
+        if max_tokens <= 0:
+            # Unlimited, depending on n_ctx.
+            max_tokens = llama_cpp.llama_n_ctx(self.ctx) - len(prompt_tokens)
+
         # Truncate max_tokens if requested tokens would exceed the context window
         max_tokens = (
             max_tokens
