Create inference.py

inference.py

@@ -0,0 +1,32 @@
+from transformers import AutoTokenizer, AutoModelForCausalLM
+import torch
+
+MODEL_NAME = "distilgpt2"
+
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)
+
+with open("prompt.txt", "r") as f:
+    SYSTEM_PROMPT = f.read().strip()
+
+def generate(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(
+            **inputs,
+            max_new_tokens=150,
+            temperature=0.7,
+            top_p=0.9,
+            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))