Update app.py

app.py

@@ -1,40 +1,94 @@
-import torch
-from transformers import AutoTokenizer, AutoModelForCausalLM
 from flask import Flask, request, jsonify
-from threading import Thread
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
+import os
 
 app = Flask(__name__)
 
-tokenizer = AutoTokenizer.from_pretrained("HuggingFaceH4/zephyr-7b-beta")
-model = AutoModelForCausalLM.from_pretrained(
-    "HuggingFaceH4/zephyr-7b-beta",
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
-
-@app.route("/api/chat", methods=["POST"])
-def chat():
-    data = request.get_json()
-    question = data.get("question", "")
-
-    prompt = f"""Eres BITER, un mentor experto en negocios con mentalidad de CEO. Respondes SIEMPRE en español y ayudas a emprendedores a tomar decisiones rápidas, inteligentes y estratégicas.
-
-Tu estilo es directo, profesional y humano. Tus respuestas son claras, realistas, y con visión práctica. Nunca usas tecnicismos innecesarios. A veces puedes ser exigente si la idea no está bien pensada, pero siempre propones formas de mejorarla.
+# 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
 
-Solo respondes a la última pregunta del usuario.
+# 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!")
 
-Usuario: {question}
-BITER:"""
+# 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})
     
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(**inputs, max_new_tokens=200)
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    respuesta_final = response.split("BITER:")[-1].strip()
-
-    return jsonify({"choices": [{"message": {"content": respuesta_final}}]})
+    except Exception as e:
+        print(f"Error: {str(e)}")
+        return jsonify({"error": f"Error al generar respuesta: {str(e)}"}), 500
 
-def run():
-    app.run(host='0.0.0.0', port=7860)
+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)
 
-Thread(target=run).start()