Update README.md

README.md

@@ -6,24 +6,19 @@ datasets:
 - zjkarina/Vikhr_instruct
 ---
 ```python
+from peft import PeftModel, PeftConfig
 from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig
+import torch
 
 MODEL_NAME = "Vikhrmodels/Vikhr-7B-instruct"
-TEMPLATE = "<s>{role}\n{content}</s>\n"
-SYSTEM_PROMPT = "Ты – полезный помощник по имени Вихрь. Ты разговариваешь с людьми и помогаешь им."
-
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
-model.to('cuda')
-model.eval()
-
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
-generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
+DEFAULT_MESSAGE_TEMPLATE = "<s>{role}\n{content}</s>\n"
+DEFAULT_SYSTEM_PROMPT = "Ты — Вихрь, русскоязычный автоматический ассистент. Ты разговариваешь с людьми и помогаешь им."
 
 class Conversation:
     def __init__(
         self,
-        message_template=TEMPLATE,
-        system_prompt=SYSTEM_PROMPT,
+        message_template=DEFAULT_MESSAGE_TEMPLATE,
+        system_prompt=DEFAULT_SYSTEM_PROMPT,
     ):
         self.message_template = message_template
         self.messages = [{
@@ -56,6 +51,24 @@ def generate(model, tokenizer, prompt, generation_config):
     output_ids = output_ids[len(data["input_ids"][0]):]
     output = tokenizer.decode(output_ids, skip_special_tokens=True)
     return output.strip()
+    
+config = PeftConfig.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(
+    config.base_model_name_or_path,
+    load_in_8bit=True,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
+model = PeftModel.from_pretrained(
+    model,
+    MODEL_NAME,
+    torch_dtype=torch.float16
+)
+model.eval()
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=False)
+generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
+print(generation_config)
 
 inputs = ["Как тебя зовут?", "Кто такой Колмогоров?"]