File size: 2,600 Bytes
acceed5 d3f31fd |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 |
from flask import Flask, request, jsonify
from flask_cors import CORS
from transformers import pipeline
import os
app = Flask(__name__)
CORS(app) # Habilita CORS para permitir peticiones desde tu frontend
# Nombre del modelo de Hugging Face
MODEL_ID = "sarvamai/sarvam-m"
# Variable para almacenar el pipeline de la IA
pipe = None
# Función para cargar el modelo
def load_model():
global pipe
try:
print(f"Cargando modelo '{MODEL_ID}'. Esto puede tardar unos minutos...")
# El token de Hugging Face se obtiene automáticamente de la variable de entorno HF_TOKEN
# que configuras como un "Secret" en Hugging Face Spaces.
pipe = pipeline("text-generation", model=MODEL_ID, device=-1) # device=-1 usa CPU
print(f"Modelo '{MODEL_ID}' cargado con éxito.")
except Exception as e:
print(f"Error al cargar el modelo '{MODEL_ID}': {e}")
pipe = None
print("El servidor no podrá responder a las solicitudes de IA.")
# Cargar el modelo inmediatamente después de crear la aplicación Flask
# Esto se asegura de que el modelo se cargue al inicio del servidor
with app.app_context():
load_model()
@app.route('/chat', methods=['POST'])
def chat():
if pipe is None:
return jsonify({"error": "El modelo de IA no está cargado. Por favor, revisa el log del servidor."}), 503
data = request.json
messages = data.get('messages') # El frontend envía el historial completo
if not messages:
return jsonify({"error": "No se proporcionaron mensajes."}), 400
try:
prompt = ""
for msg in messages:
if msg['role'] == 'user':
prompt += f"Usuario: {msg['content']}\n"
elif msg['role'] == 'assistant':
prompt += f"Asistente: {msg['content']}\n"
prompt += "Asistente:"
response = pipe(prompt, max_new_tokens=250, temperature=0.7, do_sample=True, clean_up_tokenization_spaces=True)
generated_text = response[0]['generated_text']
ai_response_content = generated_text.split("Asistente:", 1)[-1].strip()
if not ai_response_content:
ai_response_content = "Lo siento, no pude generar una respuesta clara. ¿Puedes intentarlo de nuevo?"
return jsonify({"response": ai_response_content})
except Exception as e:
print(f"Error en la inferencia de la IA: {e}")
return jsonify({"error": f"Error interno del servidor: {str(e)}"}), 500
if __name__ == '__main__':
port = int(os.environ.get('PORT', 7860))
app.run(debug=False, host='0.0.0.0', port=port) |