Update inference.py

inference.py

@@ -1,25 +1,34 @@
 import json
-from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-with open("config.json", "r") as f:
+with open("config.json") as f:
     cfg = json.load(f)
 
-BASE_MODEL = cfg["base_model"]
-
-tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
-model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
+tokenizer = AutoTokenizer.from_pretrained(cfg["base_model"])
+model = AutoModelForCausalLM.from_pretrained(
+    cfg["base_model"],
+    torch_dtype=torch.float32,
+    device_map="cpu"
+)
 
-with open("prompt.txt", "r") as f:
+with open("prompt.txt") as f:
     SYSTEM_PROMPT = f.read().strip()
 
 def chat(user_input):
-    prompt = f"{SYSTEM_PROMPT}\n\nUser: {user_input}\nBrad AI:"
-    inputs = tokenizer(prompt, return_tensors="pt")
+    messages = [
+        {"role": "system", "content": SYSTEM_PROMPT},
+        {"role": "user", "content": user_input}
+    ]
+
+    input_ids = tokenizer.apply_chat_template(
+        messages,
+        return_tensors="pt"
+    )
 
     with torch.no_grad():
         output = model.generate(
-            **inputs,
+            input_ids,
             max_new_tokens=cfg["max_new_tokens"],
             temperature=cfg["temperature"],
             top_p=cfg["top_p"],
@@ -27,3 +36,10 @@ def chat(user_input):
         )
 
     return tokenizer.decode(output[0], skip_special_tokens=True)
+
+if __name__ == "__main__":
+    while True:
+        msg = input("You: ")
+        if msg.lower() in ("exit", "quit"):
+            break
+        print(chat(msg))