Update inference.py

inference.py

@@ -1,32 +1,29 @@
+import json
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
-MODEL_NAME = "distilgpt2"
+with open("config.json", "r") as f:
+    cfg = json.load(f)
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+BASE_MODEL = cfg["base_model"]
+
+tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL)
+model = AutoModelForCausalLM.from_pretrained(BASE_MODEL)
 
 with open("prompt.txt", "r") as f:
     SYSTEM_PROMPT = f.read().strip()
 
-def generate(user_input):
+def chat(user_input):
     prompt = f"{SYSTEM_PROMPT}\n\nUser: {user_input}\nBrad AI:"
     inputs = tokenizer(prompt, return_tensors="pt")
 
     with torch.no_grad():
-        outputs = model.generate(
+        output = model.generate(
             **inputs,
-            max_new_tokens=150,
-            temperature=0.7,
-            top_p=0.9,
+            max_new_tokens=cfg["max_new_tokens"],
+            temperature=cfg["temperature"],
+            top_p=cfg["top_p"],
             do_sample=True
         )
 
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-if __name__ == "__main__":
-    while True:
-        user = input("You: ")
-        if user.lower() in ["exit", "quit"]:
-            break
-        print(generate(user))
+    return tokenizer.decode(output[0], skip_special_tokens=True)