Update README.md

README.md

@@ -28,9 +28,27 @@ tokenizer = GPT2Tokenizer.from_pretrained("RaushanTurganbay/GPT2_instruct_tuned"
 model = GPT2LMHeadModel.from_pretrained("RaushanTurganbay/GPT2_instruct_tuned")
 
 # Generate responses
-prompt = "Your input prompt here"
+class StoppingCriteriaSub(StoppingCriteria):
+    def __init__(self, stops=[], encounters=1):
+        super().__init__()
+        self.stops = [stop.to("cuda") for stop in stops]
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor):
+        for stop in self.stops:
+            if torch.all((stop == input_ids[0][-len(stop):])).item():
+                return True
+        return False
+
+
+def stopping_criteria(tokenizer, stop_words):
+    stop_words_ids = [tokenizer(stop_word, return_tensors='pt')['input_ids'].squeeze() for stop_word in stop_words]
+    stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
+    return stopping_criteria
+
+# Generate responses
+stopping = stopping_criteria(tokenizer, ["\n\nHuman:"])
+prompt = "\n\nHuman: {your_instruction}\n\nAssistant:"
 inputs = tokenizer(prompt, return_tensors="pt")
-outputs = model.generate(**inputs, max_length=150, num_return_sequences=1)
+outputs = model.generate(**inputs,  stopping_criteria=stopping, max_length=150)
 
-print(tokenizer.decode(outputs[0], skip_special_tokens=True))
+print("Model Response:", tokenizer.batch_decode(outputs))
 ```