Update README.md

README.md

@@ -43,19 +43,46 @@ This is a specialized AI assistant for programming in **BlackBox Component Build
 ## How to Use
 
 ```python
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import BitsAndBytesConfig, AutoModelForCausalLM, AutoTokenizer
+import torch
+from peft import PeftModel
 
-model_name = "hodza/BlackBox-Coder-3B"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+assert torch.cuda.is_available(), "you need cuda for this part"
+device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
 
-def get_assistant_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt")
+base_model_name = 'Qwen/Qwen2.5-Coder-3B-Instruct'
+qlora_adapter = "hodza/BlackBox-Coder-3B"
+tokenizer = AutoTokenizer.from_pretrained(base_model_name)
+
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_use_double_quant=True,
+    bnb_4bit_compute_dtype=torch.bfloat16
+)
+base_model = AutoModelForCausalLM.from_pretrained(base_model_name, device_map=device,quantization_config=bnb_config,)
+
+model = PeftModel.from_pretrained(base_model, qlora_adapter, device_map=device)
+# Define the chat template
+def format_chat_prompt(user_query):
+    return [
+        {"role": "system", "content": "You are a helpful coding assistant for BlackBox Component Builder using Component Pascal."},
+        {"role": "user", "content": user_query}
+    ]
+
+def get_assistant_response(user_query):
+    # Format the prompt using the chat template
+    chat_prompt = format_chat_prompt(user_query)
+    inputs = tokenizer.apply_chat_template(chat_prompt, return_tensors="pt").to(model.device)
+    
+    # Generate the response
     outputs = model.generate(
-        inputs.input_ids,
+        inputs,
         max_new_tokens=256,
-        temperature=0.7,
-        top_p=0.9,
+        temperature=0.3,
+        top_p=0.3,
         pad_token_id=tokenizer.eos_token_id
     )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+print(get_assistant_response("Как мне вывести массив в Log?"))