from flask import Flask, request, jsonify, render_template_string from vllm import LLM, SamplingParams from langchain_community.cache import GPTCache app = Flask(__name__) Configuración de los modelos de lenguaje modelos = { "facebook/opt-125m": LLM(model="facebook/opt-125m"), "llama-3.2-1B": LLM(model="Hjgugugjhuhjggg/llama-3.2-1B-spinquant-hf"), "gpt2": LLM(model="gpt2") } 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 = """ Documentación de la API

API de Generación de Texto

Endpoints

""" html_code_chatbot = """ Chatbot

Chatbot

""" @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__': app.run(host='0.0.0.0', port=7860)