Update inference.py

inference.py

@@ -0,0 +1,41 @@
+import torch
+from transformers import AutoTokenizer, OpenAIGPTLMHeadModel
+from evo_model import EvoTransformerV22
+
+# Load Evo model
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+evo_model = EvoTransformerV22()
+evo_model.load_state_dict(torch.load("trained_model/evo_hellaswag.pt", map_location=device))
+evo_model.to(device)
+evo_model.eval()
+
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+
+# 🧠 Evo logic
+def get_evo_response(query, context):
+    combined = query + " " + context
+    inputs = tokenizer(combined, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
+    input_ids = inputs["input_ids"].to(device)
+
+    with torch.no_grad():
+        logits = evo_model(input_ids)
+        pred = torch.argmax(logits, dim=1).item()
+
+    return f"Evo suggests: Option {pred + 1}"  # Assumes binary classification (0 or 1)
+
+# 🤖 GPT-3.5 comparison (optional)
+import openai
+openai.api_key = "sk-..."  # Replace with your OpenAI API key
+
+def get_gpt_response(query, context):
+    try:
+        prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": prompt}],
+            temperature=0.3
+        )
+        return response['choices'][0]['message']['content'].strip()
+    except Exception as e:
+        return f"Error from GPT: {e}"