from transformers import AutoModelForCausalLM, AutoTokenizer
from flask import Flask, request, jsonify
import torch

app = Flask(__name__)

# 1️⃣ Cargar el modelo GPT-J español 📌
model_name = "EleutherAI/gpt-j-6B"
print("Cargando modelo...")

# Usamos accelerate para manejar el uso de GPU/CPU
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",     # usa GPU si está disponible
    torch_dtype=torch.float16,  # optimiza memoria
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

print("Modelo cargado, listo para conversar 🤖")

# 2️⃣ Función para generar respuesta simulando chat 💬
def chat_response(prompt, history):
    full_prompt = ""
    for i, turn in enumerate(history):
        full_prompt += f"Usuario: {turn['user']}\nAsistente: {turn['bot']}\n"
    full_prompt += f"Usuario: {prompt}\nAsistente:"

    inputs = tokenizer(full_prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(
        **inputs,
        max_new_tokens=150,
        temperature=0.7,
        do_sample=True,
        top_p=0.9
    )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response.split("Asistente:")[-1].strip()

# 3️⃣ API para conversar 🧠
@app.route("/chat", methods=["POST"])
def chat():
    data = request.get_json()
    prompt = data.get("prompt", "")
    history = data.get("history", [])
    response = chat_response(prompt, history)
    return jsonify({"response": response})

# 4️⃣ Ejecutar el servidor
if __name__ == "__main__":
    print("🚀 Servidor iniciado en http://localhost:8000/chat")
    app.run(host="0.0.0.0", port=8000)