timpal0l commited on
Commit
eff2613
·
verified ·
1 Parent(s): 67629f0

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +18 -3
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
- text_to_translate = "I like to eat ice cream in the summer."
29
- prompt = f"<|endoftext|><s>User: Översätt till Svenska från Engelska\n{text_to_translate}<s>Bot:"
 
 
 
 
 
 
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])