Update README.md

README.md

@@ -75,9 +75,11 @@ Kul att du tycker det!
 ...
 ```
 
-The procedure to generate text is the same as before:
+The procedure to generate text in chat format:
 
 ```python
+from transformers import StoppingCriteriaList, StoppingCriteria
+
 prompt = """
 <|endoftext|><s>
 User:
@@ -86,17 +88,28 @@ Varför är träd fina?
 Bot:
 """.strip()
 
+# (Optional) - define a stopping criteria
+# We ideally want the model to stop generate once the response from the Bot is generated
+class StopOnTokenCriteria(StoppingCriteria):
+    def __init__(self, stop_token_id):
+        self.stop_token_id = stop_token_id
+
+    def __call__(self, input_ids, scores, **kwargs):
+        return input_ids[0, -1] == self.stop_token_id
+
+stop_on_token_criteria = StopOnTokenCriteria(stop_token_id=tokenizer.bos_token_id)
 input_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(device)
 
 generated_token_ids = model.generate(
     inputs=input_ids,
-    max_new_tokens=100,
+    max_new_tokens=128,
     do_sample=True,
     temperature=0.6,
     top_p=1,
+    stopping_criteria=StoppingCriteriaList([stop_on_token_criteria])
 )[0]
 
-generated_text = tokenizer.decode(generated_token_ids)    
+generated_text = tokenizer.decode(generated_token_ids[len(input_ids[0]):-1])
 ```
 
 Generating text using the `generate` method is done as follows:
@@ -171,7 +184,7 @@ Following Mitchell et al. (2018), we provide a model card for GPT-SW3.
 
 - Conversational
     - Familjeliv (https://www.familjeliv.se/)
-    - Flashback (https://flashback.se/)
+    - Flashback (https://flashback.org/)
     - Datasets collected through Parlai (see Appendix in data paper for complete list) (https://github.com/facebookresearch/ParlAI)
     - Pushshift.io Reddit dataset, developed in Baumgartner et al. (2020) and processed in Roller et al. (2021)