Update inference.py

inference.py

@@ -1,25 +1,45 @@
-from transformers import AutoTokenizer
+import openai
+from transformers import BertTokenizer
 from evo_model import EvoTransformerForClassification
-from init_save import initialize_and_save_model  # Ensure this line is added
 import torch
 
-# ✅ Ensure model is initialized and saved BEFORE loading
-initialize_and_save_model("trained_model")
-
-# 🔁 Load tokenizer and model
-tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
+# Load Evo model
 model = EvoTransformerForClassification.from_pretrained("trained_model")
 model.eval()
 
-def generate_response(goal, sol1, sol2):
-    prompt = f"Goal: {goal}\nOption 1: {sol1}\nOption 2: {sol2}\nWhich is better?"
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
+# Tokenizer (BERT-compatible)
+tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+
+# Set OpenAI key (assumes you have it set as ENV VAR or replace directly)
+import os
+openai.api_key = os.getenv("OPENAI_API_KEY")
 
-    if 'token_type_ids' in inputs:
-        del inputs['token_type_ids']  # Evo doesn't use this
+def query_gpt35(prompt):
+    try:
+        response = openai.ChatCompletion.create(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": prompt}],
+            max_tokens=50,
+            temperature=0.3,
+        )
+        return response['choices'][0]['message']['content'].strip()
+    except Exception as e:
+        return f"[GPT-3.5 Error] {e}"
 
+def generate_response(goal, option1, option2):
+    # Evo prediction
+    prompt = f"Goal: {goal}\nOption 1: {option1}\nOption 2: {option2}\nWhich is better?"
+    inputs = tokenizer([goal + " " + option1, goal + " " + option2],
+                       return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
         logits = model(**inputs)
-    
-    predicted = torch.argmax(logits, dim=1).item()
-    return f"Option {predicted + 1} seems more reasonable based on EvoTransformer."
+    pred = torch.argmax(logits, dim=1).item()
+    evo_result = option1 if pred == 0 else option2
+
+    # GPT-3.5 prediction
+    gpt_result = query_gpt35(prompt)
+
+    return {
+        "evo_suggestion": evo_result,
+        "gpt_suggestion": gpt_result
+    }