Update README.md

README.md

@@ -36,7 +36,7 @@ I trained with 1x A6000 GPUs for 90 hours.
 ```python
 from transformers import AutoModelForCausalLM, AutoTokenizer, StoppingCriteria, StoppingCriteriaList
 import torch
-repo = "squarelike/Gugugo-koen-7B-V1.1-GPTQ"
+repo = "squarelike/Gugugo-koen-7B-V1.1"
 model = AutoModelForCausalLM.from_pretrained(
         repo,
         load_in_4bit=True
@@ -56,7 +56,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=""):
@@ -69,13 +69,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("</끝>", "")