Update inference.py

inference.py

@@ -7,11 +7,23 @@ import os
 
 # Load Evo model and tokenizer
 model = EvoTransformerV22()
-model.load_state_dict(torch.load("evo_hellaswag.pt", map_location="cpu"))
-model.eval()
-
 tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
 
+# Smart load order
+if os.path.exists("trained_model/evo_retrained.pt"):
+    model.load_state_dict(torch.load("trained_model/evo_retrained.pt", map_location="cpu"))
+    print("🔁 Loaded retrained Evo model.")
+elif os.path.exists("trained_model/evo_pretrained.pt"):
+    model.load_state_dict(torch.load("trained_model/evo_pretrained.pt", map_location="cpu"))
+    print("📦 Loaded pretrained Evo model.")
+elif os.path.exists("evo_hellaswag.pt"):
+    model.load_state_dict(torch.load("evo_hellaswag.pt", map_location="cpu"))
+    print("📥 Loaded default Evo model.")
+else:
+    raise FileNotFoundError("❌ No Evo model file found.")
+
+model.eval()
+
 # GPT Setup
 openai.api_key = os.getenv("OPENAI_API_KEY")  # 🔒 Load securely from environment