Needed for latest transformers

nuprl · Feb 25, 2024 · 8b683a9 · 8b683a9
1 parent 3609be3
commit 8b683a9
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/automodel.py b/automodel.py
@@ -45,9 +45,18 @@ def completion_tensors(
         temperature: float,
         top_p: float,
     ):
+        self.model.eval() # Not essential, but just in case.
+
         inputs = self.tokenizer(
-            prompts, padding=True, return_tensors="pt", return_token_type_ids=False
+            prompts,
+            padding=True,
+            return_tensors="pt",
+            return_token_type_ids=False,
+            truncation=True,
+            max_length=max_length - 1,
         ).to("cuda")
+
+
         with torch.no_grad():
             output = self.model.generate(
                 **inputs,