Update README.md

README.md

@@ -54,7 +54,7 @@ class StoppingCriteriaSub(StoppingCriteria):
 
         return False
 
-stop_words_ids = torch.tensor([[829, 45107, 29958], [1533, 45107, 29958], [829, 45107, 29958], [21106, 45107, 29958]])
+stop_words_ids = torch.tensor([[829, 45107, 29958], [1533, 45107, 29958], [829, 45107, 29958], [21106, 45107, 29958]]).to("cuda")
 stopping_criteria = StoppingCriteriaList([StoppingCriteriaSub(stops=stop_words_ids)])
 
 def gen(lan="en", x=""):
@@ -67,13 +67,10 @@ def gen(lan="en", x=""):
             prompt,
             return_tensors='pt',
             return_token_type_ids=False
-        ),
-        max_new_tokens=1000,
+        ).to("cuda"),
+        max_new_tokens=2000,
         temperature=0.1,
-        no_repeat_ngram_size=10,
-        early_stopping=True,
         do_sample=True,
-        eos_token_id=2,
         stopping_criteria=stopping_criteria
     )
     return tokenizer.decode(gened[0][1:]).replace(prompt+" ", "").replace("</끝>", "")