Update README.md

README.md

@@ -58,18 +58,44 @@ prompt += '''Instruction:\tYou are to try to answer the following question using
 Instruction:\tYour response should be a well formed JSON object with an 'answerable' property followed by an 'answer' property.
 Instruction:\tIf you cannot answer the question given the information, the value of the 'answerable' should be 'false' and the 'answer' should be an empty string.
 Instruction:\tIf you can answer the question given the information, the value of the 'answerable' should be 'true' and your answer should be the string value of the 'answer' property.
-''' + info + qs + " Response:"
+''' + info + qs
 ```
 
 Inference:
 
+We recommend using newline character for stopping criterion, as follows:
+
+```python
+from transformers import StoppingCriteria, StoppingCriteriaList
+
+eos_tokens = [tokenizer.eos_token,'\n']
+eos_token_ids = [tokenizer.encode(token)[0] for token in eos_tokens]
+
+class MultipleEOSTokensStoppingCriteria(StoppingCriteria):
+    def __init__(self, eos_token_ids):
+        self.eos_token_ids = set(eos_token_ids)
+    def __call__(self, input_ids, scores) -> bool:
+        if input_ids.shape[-1] <= 1:
+            return False
+        for eos_token_id in self.eos_token_ids:
+            if eos_token_id == input_ids[0, -1].item():
+                return True
+        return False
+
+# Define stopping criteria
+multiple_eos_tokens_processor = MultipleEOSTokensStoppingCriteria(eos_token_ids)
+stopping_criteria = StoppingCriteriaList([multiple_eos_tokens_processor])
+```
+It can be used in inference as follows:
+
 ```python
 inputs = tokenizer(prompt, return_tensors="pt").to(device)
 generate_ids = model.generate(
     **inputs,
     max_new_tokens=1024,
     temperature=0.0,
-    num_beams=2
+    num_beams=2,
+    stopping_criteria=stopping_criteria
 )
 response = tokenizer.decode(generate_ids[0],
                       skip_special_tokens=True,