Update run_model.py

run_model.py

@@ -1,4 +1,4 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 from huggingface_hub import hf_hub_download
 import os
 import torch
@@ -11,23 +11,26 @@ SYSTEM_PROMPT = """You are helpful AI assistant. You answer questions truthfully
 
 MODEL_NAME = "microsoft/Phi-3.5-mini-instruct"
 
+quantization_config = BitsAndBytesConfig(load_in_4bit=True)
+
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME, 
     device_map="cpu", 
-    torch_dtype=torch.bfloat16,
+    torch_dtype="auto",
+    quantization_config=quantization_config
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
 # print(f"Microsoft Phi-3.5-mini-instruct Downloaded Successfully !")
 
 
-generation_args = {
-    "max_new_tokens": 64,
-    "return_full_text": False,
-    "temperature": 0.1,
-    "top_p": 1.0,
-    "do_sample": False,
-}
+# generation_args = {
+#     "max_new_tokens": 64,
+#     "return_full_text": False,
+#     "temperature": 0.1,
+#     "top_p": 1.0,
+#     "do_sample": True,
+# }
 
 # pipe = pipeline(
 #     "text-generation",
@@ -49,7 +52,7 @@ async def generate_response(prompt: str, context: str = "", history: list = [])
     with torch.no_grad():
         prompt_text = "\n".join([f"{m['role']}: {m['content']}" for m in message])
         inputs = tokenizer(prompt_text, return_tensors="pt").to(model.device)
-        outputs = model.generate(**inputs, **generation_args) 
+        outputs = model.generate(**inputs, temperature=0.1, do_resample=True) 
         outputs = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
     # return outputs[0]["generated_text"]