| from flask import Flask, request, jsonify, render_template_string |
| from vllm import LLM, SamplingParams |
| from langchain_community.cache import GPTCache |
| import torch |
|
|
| app = Flask(__name__) |
|
|
| |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
| |
| try: |
| modelos = { |
| "facebook/opt-125m": LLM(model="facebook/opt-125m", device=device), |
| "llama-3.2-1B": LLM(model="Hjgugugjhuhjggg/llama-3.2-1B-spinquant-hf", device=device), |
| "gpt2": LLM(model="gpt2", device=device) |
| } |
| except KeyError as e: |
| print(f"Error al inicializar el modelo con {device}: {e}") |
| modelos = {} |
|
|
| |
| if not modelos: |
| print("Error: No se pudo inicializar ningún modelo.") |
| exit(1) |
|
|
| |
| caches = { |
| nombre: GPTCache(modelo, max_size=1000) |
| for nombre, modelo in modelos.items() |
| } |
|
|
| |
| sampling_params = SamplingParams(temperature=0.8, top_p=0.95) |
|
|
| |
| html_code_docs = """ |
| <!DOCTYPE html> |
| <html> |
| <head> |
| <title>Documentación de la API</title> |
| </head> |
| <body> |
| <h1>API de Generación de Texto</h1> |
| <h2>Endpoints</h2> |
| <ul> |
| <li> |
| <h3>Generar texto</h3> |
| <p>Método: POST</p> |
| <p>Ruta: /generate</p> |
| <p>Parámetros:</p> |
| <ul> |
| <li>prompts: Lista de prompts para generar texto</li> |
| <li>modelo: Nombre del modelo a utilizar</li> |
| </ul> |
| <p>Ejemplo:</p> |
| <pre>curl -X POST -H "Content-Type: application/json" -d '{"prompts": ["Hola, cómo estás?"], "modelo": "facebook/opt-125m"}' http://localhost:5000/generate</pre> |
| </li> |
| <li> |
| <h3>Obtener lista de modelos</h3> |
| <p>Método: GET</p> |
| <p>Ruta: /modelos</p> |
| <p>Ejemplo:</p> |
| <pre>curl -X GET http://localhost:5000/modelos</pre> |
| </li> |
| <li> |
| <h3>Chatbot</h3> |
| <p>Método: POST</p> |
| <p>Ruta: /chatbot</p> |
| <p>Parámetros:</p> |
| <ul> |
| <li>mensaje: Mensaje para el chatbot</li> |
| <li>modelo: Nombre del modelo a utilizar</li> |
| </ul> |
| <p>Ejemplo:</p> |
| <pre>curl -X POST -H "Content-Type: application/json" -d '{"mensaje": "Hola, cómo estás?", "modelo": "facebook/opt-125m"}' http://localhost:5000/chatbot</pre> |
| </li> |
| </ul> |
| </body> |
| </html> |
| """ |
|
|
| |
| html_code_chatbot = """ |
| <!DOCTYPE html> |
| <html> |
| <head> |
| <title>Chatbot</title> |
| </head> |
| <body> |
| <h1>Chatbot</h1> |
| <form id="chat-form"> |
| <input type="text" id="mensaje" placeholder="Escribe un mensaje"> |
| <button type="submit">Enviar</button> |
| </form> |
| <div id="respuestas"></div> |
| |
| <script> |
| const form = document.getElementById('chat-form'); |
| const mensajeInput = document.getElementById('mensaje'); |
| const respuestasDiv = document.getElementById('respuestas'); |
| |
| form.addEventListener('submit', (e) => { |
| e.preventDefault(); |
| const mensaje = mensajeInput.value; |
| fetch('/chatbot', { |
| method: 'POST', |
| headers: { |
| 'Content-Type': 'application/json' |
| }, |
| body: JSON.stringify({ mensaje }) |
| }) |
| .then((res) => res.json()) |
| .then((data) => { |
| const respuesta = data.respuesta; |
| const respuestaHTML = `<p>Tú: ${mensaje}</p><p>Chatbot: ${respuesta}</p>`; |
| respuestasDiv.innerHTML += respuestaHTML; |
| mensajeInput.value = ''; |
| }); |
| }); |
| </script> |
| </body> |
| </html> |
| """ |
|
|
| @app.route('/generate', methods=['POST']) |
| def generate(): |
| data = request.get_json() |
| prompts = data.get('prompts', []) |
| modelo_seleccionado = data.get('modelo', "facebook/opt-125m") |
|
|
| if modelo_seleccionado not in modelos: |
| return jsonify({"error": "Modelo no encontrado"}), 404 |
|
|
| outputs = caches[modelo_seleccionado].generate(prompts, sampling_params) |
|
|
| results = [] |
| for output in outputs: |
| prompt = output.prompt |
| generated_text = output.outputs[0].text |
| results.append({ |
| 'prompt': prompt, |
| 'generated_text': generated_text |
| }) |
|
|
| return jsonify(results) |
|
|
| @app.route('/modelos', methods=['GET']) |
| def get_modelos(): |
| return jsonify({"modelos": list(modelos.keys())}) |
|
|
| @app.route('/docs', methods=['GET']) |
| def docs(): |
| return render_template_string(html_code_docs) |
|
|
| @app.route('/chatbot', methods=['POST']) |
| def chatbot(): |
| data = request.get_json() |
| mensaje = data.get('mensaje', '') |
| modelo_seleccionado = data.get('modelo', "facebook/opt-125m") |
|
|
| if modelo_seleccionado not in modelos: |
| return jsonify({"error": "Modelo no encontrado"}), 404 |
|
|
| outputs = caches[modelo_seleccionado].generate([mensaje], sampling_params) |
|
|
| respuesta = outputs[0].outputs[0].text |
|
|
| return jsonify({"respuesta": respuesta}) |
|
|
| @app.route('/chat', methods=['GET']) |
| def chat(): |
| return render_template_string(html_code_chatbot) |
|
|
| if __name__ == '__main__': |
| |
| if modelos: |
| app.run(host='0.0.0.0', port=7860) |
| else: |
| print("Error: No se pudieron cargar los modelos. El servidor no se iniciará.") |
|
|
