Update README.md

README.md

@@ -88,7 +88,7 @@ Researchers should apply responsible AI best practices, including mapping, measu
 ### Loading the model locally
 
     import torch
-    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, StoppingCriteria
 
     torch.random.manual_seed(0)
 
@@ -114,11 +114,20 @@ Researchers should apply responsible AI best practices, including mapping, measu
         tokenizer=tokenizer,
     )
 
+    class EosListStoppingCriteria(StoppingCriteria):
+      def __init__(self, eos_sequence = [32007]):
+          self.eos_sequence = eos_sequence
+  
+      def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+          last_ids = input_ids[:,-len(self.eos_sequence):].tolist()
+          return self.eos_sequence in last_ids
+
     generation_args = {
         "max_new_tokens": 500,
         "return_full_text": False,
         "temperature": 0.0,
         "do_sample": False,
+        "generation_kwargs": {"stopping_criteria": EosListStoppingCriteria()}
     }
 
     output = pipe(messages, **generation_args)