import torch
from transformers import AutoTokenizer, OpenAIGPTLMHeadModel
from evo_model import EvoTransformerV22

# 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
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 = torch.argmax(logits, dim=1).item()

    return f"Evo suggests: Option {pred + 1}"  # Assumes binary classification (0 or 1)

# 🤖 GPT-3.5 comparison (optional)
import openai
openai.api_key = "sk-..."  # Replace with your OpenAI API key

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