inference.py

import torch
from transformers import AutoTokenizer
from evo_model import EvoTransformerV22  ✅ CORRECT
from search_utils import web_search
import openai

# 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")

# GPT Setup
openai.api_key = "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"  # 🔑 Set your actual key securely

def get_evo_response(query, options, user_context=""):
    context_texts = web_search(query) + ([user_context] if user_context else [])
    context_str = "\n".join(context_texts)
    input_pairs = [f"{query} [SEP] {opt} [CTX] {context_str}" for opt in options]

    scores = []
    for pair in input_pairs:
        encoded = tokenizer(pair, return_tensors="pt", truncation=True, padding="max_length", max_length=128)
        with torch.no_grad():
            output = model(encoded["input_ids"])
            score = torch.sigmoid(output).item()
            scores.append(score)

    best_idx = int(scores[1] > scores[0])
    return (
        options[best_idx],
        f"{options[0]}: {scores[0]:.3f} vs {options[1]}: {scores[1]:.3f}",
        max(scores),
        context_str
    )

def get_gpt_response(query, user_context=""):
    try:
        context_block = f"\n\nContext:\n{user_context}" if user_context else ""
        completion = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": query + context_block}]
        )
        return completion.choices[0].message.content.strip()
    except Exception as e:
        return f"⚠️ GPT error: {str(e)}"