Update README.md

README.md

@@ -19,13 +19,9 @@ GITHUB_ACTIONS=true pip install auto-gptq
 from transformers import AutoTokenizer
 from auto_gptq import AutoGPTQForCausalLM
 
-device = "cuda:0"
-quantized_model_dir = "saiga2-13b-4bit"
-tokenizer = AutoTokenizer.from_pretrained(quantized_model_dir, use_fast=True)
-model = AutoGPTQForCausalLM.from_quantized(quantized_model_dir, device=device)
 
-def generate_answer(model, tokenizer, request: str):
-    s = f"system\n{config.system_prompt}</s>\n" + \
+def generate_answer(model, tokenizer, request: str, system_prompt: str):
+    s = f"system\n{system_prompt}</s>\n" + \
         f"<s>user\n{request}</s>\n" + \
         f"<s>bot\n"
     request_tokens = tokenizer(s, return_tensors="pt")    
@@ -40,22 +36,20 @@ def generate_answer(model, tokenizer, request: str):
                     no_repeat_ngram_size=15,
                     max_new_tokens=1536,
                     pad_token_id=tokenizer.eos_token_id)[0]
-    print(request)
     answer_tokens = answer_tokens[len(request_tokens[0]):-1]
     answer = tokenizer.decode(answer_tokens).strip()
-    print(answer)
     return answer
 
-model_name = "gurgutan/saiga2-13b-4bit"
+
+model_name = "saiga2-13b-4bit"
 system_prompt = "Ты — Сайга, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
 tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)
-model = AutoGPTQForCausalLM.from_quantized(model_name, device=device)
+model = AutoGPTQForCausalLM.from_quantized(model_name, device="cuda:0")
 model.eval()
 
 user_text = "Сочини стих, который начинается словами: Буря мглою небо кроет"
-answer_text = generate_answer(model, tokenizer, user_text)  
+answer_text = generate_answer(model, tokenizer, user_text, system_prompt)  
 print(answer_text)
-
 ```
 # Original model: [saiga2-13B-lora](https://huggingface.co/IlyaGusev/saiga2_13b_lora)
 Модель ассистента на основе LLaMA2 дообученная на русскоязычных наборах. Модель имеет 13 млрд. параметров.