Update README.md

README.md

@@ -38,25 +38,33 @@ class StopOnTokenCriteria(StoppingCriteria):
         return input_ids[0, -1] == self.stop_token_id
 
 
-stop_on_token_criteria = StopOnTokenCriteria(stop_token_id=2)
-
 pipe = pipeline(
     "text-generation",
-    "AI-Sweden-Models/gpt-sw3-6.7b-v2-translator",
+    model="AI-Sweden-Models/gpt-sw3-6.7b-v2-translator",
     device=device
 )
+
+stop_on_token_criteria = StopOnTokenCriteria(stop_token_id=pipe.tokenizer.bos_token_id)
+text = "I like to eat ice cream in the summer."
+
 # This will translate English to Swedish
 # To translate from Swedish to English the prompt would be:
 # prompt = f"<|endoftext|><s>User: Översätt till Engelska från Svenska\n{text}<s>Bot:"
-text = "I like to eat ice cream in the summer."
+
 prompt = f"<|endoftext|><s>User: Översätt till Svenska från Engelska\n{text}<s>Bot:"
 
+input_tokens = pipe.tokenizer(prompt, return_tensors="pt").input_ids.to(device)
+max_model_length = 2048
+dynamic_max_length = max_model_length - input_tokens.shape[1]
+
 response = pipe(
     prompt,
-    max_length=768,
+    max_length=dynamic_max_length,
+    truncation=True,
     stopping_criteria=StoppingCriteriaList([stop_on_token_criteria])
 )
 
+# Extract and print the generated translation
 print(response[0]["generated_text"].split("<s>Bot: ")[-1])
 ```
 ```python