import torch
from transformers import AutoTokenizer
from evo_model import EvoTransformerV22
from openai import OpenAI
import os

# 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 (binary classification with sigmoid)
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 = int(torch.sigmoid(logits).item() > 0.5)

    return f"Evo suggests: Option {pred + 1}"

# 🤖 GPT-3.5 comparison using openai>=1.0.0
openai_api_key = os.environ.get("OPENAI_API_KEY", "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")  # Replace with real key or set via HF secrets
client = OpenAI(api_key=openai_api_key)

def get_gpt_response(query, context):
    try:
        prompt = f"Context: {context}\n\nQuestion: {query}\n\nAnswer:"
        response = client.chat.completions.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}"