Update run_model.py

run_model.py

@@ -15,8 +15,8 @@ MODEL_NAME = "microsoft/Phi-3.5-mini-instruct"
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME, 
-    device_map="cpu", 
-    torch_dtype="auto",
+    device_map="auto", 
+    torch_dtype=torch.float16,
 )
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 
@@ -31,13 +31,17 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 #     "do_sample": True,
 # }
 
-# pipe = pipeline(
-#     "text-generation",
-#     model=model,
-#     tokenizer=tokenizer,
-# )
 model.eval()
 
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=64, 
+    temperature=0.4, 
+    do_sample=True
+)
+
 def generate_response(prompt: str, context: str = "", history: list = []) -> str:
     
     BASE_MESSAGE = [{"role" : "system", "content" : SYSTEM_PROMPT}] + history
@@ -48,11 +52,12 @@ def generate_response(prompt: str, context: str = "", history: list = []) -> str
         message = BASE_MESSAGE + [{"role" : "user", "content" : f"Context : {context}\nQuery : {prompt}"}]
 
     # outputs = pipe(message, **generation_args)
-    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, temperature=0.1, do_sample=True, max_new_tokens=64) 
-        outputs = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    outputs = pipe(message)
+    # 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, temperature=0.1, do_sample=True, max_new_tokens=64) 
+    #     outputs = tokenizer.decode(outputs[0], skip_special_tokens=True)
         
-    # return outputs[0]["generated_text"]
-    return outputs
+    return outputs[0]["generated_text"]
+    # return outputs