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

app = Flask(__name__)

# Configuración CORS para permitir solicitudes desde tu dominio
@app.after_request
def after_request(response):
    response.headers.add('Access-Control-Allow-Origin', 'https://justbyte.es')
    response.headers.add('Access-Control-Allow-Headers', 'Content-Type')
    response.headers.add('Access-Control-Allow-Methods', 'POST')
    return response

# Cargar el modelo y tokenizador (se carga una sola vez al iniciar)
@app.before_first_request
def load_model():
    global model, tokenizer
    
    print("Cargando modelo Zephyr-7B...")
    
    # Cargar el modelo y tokenizador
    model_name = "HuggingFaceH4/zephyr-7b-beta"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,  # Usar precisión reducida para ahorrar memoria
        device_map="auto",          # Distribuir el modelo automáticamente
        load_in_8bit=True           # Cuantización a 8 bits para reducir uso de memoria
    )
    
    print("Modelo cargado correctamente!")

# Cargar el prompt desde el archivo
def get_system_prompt():
    with open("prompt.txt", "r", encoding="utf-8") as f:
        return f.read().strip()

@app.route('/generate', methods=['POST'])
def generate_response():
    try:
        # Obtener la pregunta del usuario
        data = request.json
        user_message = data.get('message', '')
        
        if not user_message:
            return jsonify({"error": "No se proporcionó ninguna pregunta"}), 400
        
        # Obtener el prompt del sistema
        system_prompt = get_system_prompt()
        
        # Crear el formato de conversación para Zephyr
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_message}
        ]
        
        # Convertir mensajes al formato que espera el modelo
        prompt = tokenizer.apply_chat_template(messages, tokenize=False)
        
        # Generar respuesta
        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        
        # Configuración de generación
        generation_config = {
            "max_new_tokens": 500,
            "temperature": 0.7,
            "top_p": 0.9,
            "do_sample": True,
            "pad_token_id": tokenizer.eos_token_id
        }
        
        # Generar respuesta
        with torch.no_grad():
            outputs = model.generate(**inputs, **generation_config)
        
        # Decodificar la respuesta
        full_response = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # Extraer solo la respuesta del asistente (después del último mensaje del usuario)
        assistant_response = full_response.split("assistant:")[-1].strip()
        
        return jsonify({"response": assistant_response})
    
    except Exception as e:
        print(f"Error: {str(e)}")
        return jsonify({"error": f"Error al generar respuesta: {str(e)}"}), 500

if __name__ == '__main__':
    # Puerto que Hugging Face Spaces utiliza
    port = int(os.environ.get('PORT', 7860))
    app.run(host='0.0.0.0', port=port)