baigo28/embedding-api

1from flask import Flask, request, jsonify2from sentence_transformers import SentenceTransformer3import os4 5app = Flask(__name__)6 7# --- Configuración del Modelo de Embedding ---8EMBEDDING_MODEL_LOCAL = "intfloat/multilingual-e5-large"9embedding_model = None # Se cargará al iniciar la aplicación10 11def load_model():12    """Carga el modelo de SentenceTransformer al inicio de la aplicación."""13    global embedding_model14    if embedding_model is None:15        print(f"Cargando el modelo de embedding: {EMBEDDING_MODEL_LOCAL}...")16        # Usa 'cpu' si no tienes una GPU o si hay problemas con CUDA.17        # Esto cargará el modelo desde la caché si ya lo descargaste.18        embedding_model = SentenceTransformer(EMBEDDING_MODEL_LOCAL, device='cpu')19        print("Modelo de embedding cargado.")20 21# Ruta de API para generar embeddings22@app.route('/embed', methods=['POST'])23def embed_text():24    data = request.get_json()25    if not data or 'text' not in data:26        return jsonify({"error": "Se requiere un campo 'text' en el cuerpo de la solicitud."}), 40027 28    text_to_embed = data['text']29    30    try:31        # Genera el embedding32        embedding = embedding_model.encode([text_to_embed]).tolist()[0] # .tolist()[0] para obtener la lista de flotantes33        return jsonify({"embedding": embedding}), 20034    except Exception as e:35        print(f"Error al generar embedding: {e}")36        return jsonify({"error": f"Error interno al generar embedding: {str(e)}"}), 50037 38# Endpoint para verificar que el servicio está funcionando y el modelo cargado39@app.route('/health', methods=['GET'])40def health_check():41    if embedding_model:42        return jsonify({"status": "healthy", "model_loaded": True}), 20043    else:44        return jsonify({"status": "loading", "model_loaded": False}), 50345 46if __name__ == '__main__':47    # Carga el modelo cuando la aplicación Flask se inicia48    # Solo en desarrollo, en producción un servidor WSGI (Gunicorn/uWSGI) lo manejaría49    load_model() 50    # Para producción, se suele usar gunicorn: gunicorn -w 4 -b 0.0.0.0:8000 app:app51    app.run(host='0.0.0.0', port=5001)