Update README.md

README.md

@@ -59,31 +59,35 @@ To verify results locally using the transformers and peft libraries:
 ```
 
 from peft import PeftModel
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 import torch
 
-# 1. Load the base model in 4-bit for efficiency
-base_model = "mistralai/Mistral-7B-v0.3"
-model = AutoModelForCausalLM.from_pretrained(
-    base_model, 
-    load_in_4bit=True, 
-    torch_dtype=torch.float16,
+# 1. Setup 4-bit configuration
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_quant_type="nf4"
+)
+
+# 2. Load Base Model and Adapter
+base_model_id = "mistralai/Mistral-7B-v0.3"
+adapter_id = "rakshath1/it-support-mistral-7b-expert"
+
+base_model = AutoModelForCausalLM.from_pretrained(
+    base_model_id,
+    quantization_config=quantization_config,
     device_map="auto"
 )
 
-# 2. Load the Expert Adapter and Tokenizer
-adapter_path = "rakshath1/it-support-mistral-7b-expert"
-model = PeftModel.from_pretrained(model, adapter_path)
-tokenizer = AutoTokenizer.from_pretrained(adapter_path)
-tokenizer.pad_token = tokenizer.eos_token
+model = PeftModel.from_pretrained(base_model, adapter_id)
+tokenizer = AutoTokenizer.from_pretrained(base_model_id)
 
-# 3. Test a ticket
+# 3. Inference
 ticket = "### Instruction:\nTicket: 'VPN access denied for user in Mangalore office.'\n\n### Response:\n"
 inputs = tokenizer(ticket, return_tensors="pt").to("cuda")
 
-# Generate response
 with torch.no_grad():
-    outputs = model.generate(**inputs, max_new_tokens=64, pad_token_id=tokenizer.eos_token_id)
+    outputs = model.generate(**inputs, max_new_tokens=64)
 
 print(tokenizer.decode(outputs[0], skip_special_tokens=True).split("### Response:\n")[-1])