Update README.md

README.md

@@ -135,6 +135,13 @@ model = AutoModelForCausalLM.from_pretrained("norallm/normistral-11b").cuda().ev
 prompt = """Engelsk: {0}
 Bokmål:"""
 
+# Define tokens that should end the generation (any token with a newline)
+eos_token_ids = [
+    token_id
+    for token_id in range(tokenizer.vocab_size)
+    if '\n' in tokenizer.decode([token_id])
+]
+
 # Generation function
 @torch.no_grad()
 def generate(text):
@@ -144,7 +151,7 @@ def generate(text):
         input_ids,
         max_new_tokens=64,
         do_sample=False,
-        eos_token_id=tokenizer('\n').input_ids
+        eos_token_id=eos_token_ids
     )
     return tokenizer.decode(prediction[0, input_ids.size(1):]).strip()