Update README.md
Browse files
README.md
CHANGED
@@ -22,11 +22,26 @@ from transformers import pipeline, StoppingCriteriaList, StoppingCriteria
|
|
22 |
|
23 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
24 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
25 |
stop_on_token_criteria = StopOnTokenCriteria(stop_token_id=2)
|
26 |
-
pipe = pipeline("text-generation", "AI-Sweden-Models/gpt-sw3-6.7b-v2-translator", device=device)
|
27 |
|
28 |
-
|
29 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
30 |
|
31 |
response = pipe(prompt, max_length=768, stopping_criteria=StoppingCriteriaList([stop_on_token_criteria])))
|
32 |
print(response[0]["generated_text"].split("<s>Bot: ")[-1])
|
|
|
22 |
|
23 |
device = "cuda" if torch.cuda.is_available() else "cpu"
|
24 |
|
25 |
+
# (Optional) - define a stopping criteria
|
26 |
+
# We ideally want the model to stop generate once the response from the Bot is generated
|
27 |
+
class StopOnTokenCriteria(StoppingCriteria):
|
28 |
+
def __init__(self, stop_token_id):
|
29 |
+
self.stop_token_id = stop_token_id
|
30 |
+
|
31 |
+
def __call__(self, input_ids, scores, **kwargs):
|
32 |
+
return input_ids[0, -1] == self.stop_token_id
|
33 |
+
|
34 |
+
|
35 |
stop_on_token_criteria = StopOnTokenCriteria(stop_token_id=2)
|
|
|
36 |
|
37 |
+
pipe = pipeline(
|
38 |
+
"text-generation",
|
39 |
+
"AI-Sweden-Models/gpt-sw3-6.7b-v2-translator",
|
40 |
+
device=device
|
41 |
+
)
|
42 |
+
|
43 |
+
text = "I like to eat ice cream in the summer."
|
44 |
+
prompt = f"<|endoftext|><s>User: Översätt till Svenska från Engelska\n{text}<s>Bot:"
|
45 |
|
46 |
response = pipe(prompt, max_length=768, stopping_criteria=StoppingCriteriaList([stop_on_token_criteria])))
|
47 |
print(response[0]["generated_text"].split("<s>Bot: ")[-1])
|