Update inference.py

inference.py

@@ -1,24 +1,18 @@
-import torch
-from model import EvoTransformer
+from evo_model import EvoTransformerForClassification
 from transformers import AutoTokenizer
+import torch
 
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
-
-# Load model
-model = EvoTransformer(
-    vocab_size=tokenizer.vocab_size,
-    d_model=256,
-    nhead=4,
-    dim_feedforward=512,
-    num_layers=4
-)
-model.load_state_dict(torch.load("trained_model.pt", map_location=torch.device("cpu")))
+# Load model and tokenizer once
+model = EvoTransformerForClassification.from_pretrained("trained_model")
+tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 model.eval()
 
-def predict(goal, sol1, sol2):
-    text = goal + " " + sol1 + " " + sol2
-    inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=64)
+def generate_response(goal, sol1, sol2):
+    prompt = f"Goal: {goal} Option 1: {sol1} Option 2: {sol2}"
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True)
     with torch.no_grad():
-        logits = model(inputs["input_ids"])
-    return "Solution 1" if logits.argmax().item() == 0 else "Solution 2"
+        outputs = model(**inputs)
+        logits = outputs.logits if hasattr(outputs, "logits") else outputs[0]
+    prediction = torch.argmax(logits, dim=1).item()
+    return "Solution 1" if prediction == 0 else "Solution 2"
+