complete only one line

okdshin · Nov 23, 2023 · 2c5204b · 2c5204b
1 parent d700cd4
commit 2c5204b
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/flatline_lsp.py b/flatline_lsp.py
@@ -120,8 +120,8 @@ def generate_completion(self, text: str) -> str:
             tokenized_prompt = self.tokenizer(text).input_ids
             generated_tokens = self.model.generate(inputs=torch.LongTensor(
                 [tokenized_prompt]), max_new_tokens=self.max_new_tokens, do_sample=False,
-                #stopping_criteria=[stop_cutoff_completion, self.stop_word])[0]
-                stopping_criteria=[stop_cutoff_completion])[0]
+                stopping_criteria=[stop_cutoff_completion, self.stop_word])[0]
+            # stopping_criteria=[stop_cutoff_completion])[0]
             generated_text = self.tokenizer.decode(generated_tokens[len(tokenized_prompt):])
         return generated_text