import os
import logging
import asyncio
import uvicorn
import torch
import random
from transformers import AutoModelForCausalLM, AutoTokenizer
from fastapi import FastAPI, Query, HTTPException
from fastapi.responses import HTMLResponse
from typing import List
import concurrent.futures
# Configuración de logging
logging.basicConfig(level=logging.DEBUG)
logger = logging.getLogger(__name__)
# Inicializar la aplicación FastAPI
app = FastAPI()
# Diccionario para almacenar los modelos
model_dict = {}
model_lock = asyncio.Lock()
# Lista para almacenar el historial de mensajes
message_history = []
# Cargar modelos sincrónicamente al iniciar la aplicación
models_to_load = ["gpt2-medium", "gpt2-large", "gpt2"]
for model_name in models_to_load:
model = AutoModelForCausalLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model_dict[model_name] = (model, tokenizer)
logger.info(f"Successfully loaded {model_name} model")
# Ruta principal para la interfaz web
@app.get('/')
async def main():
html_code = """
ChatGPT Chatbot
"""
return HTMLResponse(content=html_code, status_code=200)
# Ruta para la generación de respuestas
@app.get('/autocomplete')
async def autocomplete(q: str = Query(...)):
global model_dict, message_history
# Generar respuestas con los tres modelos en paralelo
async def generate_response(model, tokenizer, q):
input_ids = tokenizer.encode(q, return_tensors="pt")
output = model.generate(input_ids, max_length=150, num_return_sequences=1)
response_text = tokenizer.decode(output[0], skip_special_tokens=True)
return response_text
responses = []
for model_name, (model, tokenizer) in model_dict.items():
responses.append(generate_response(model, tokenizer, q))
responses = await asyncio.gather(*responses)
# Filtrar respuestas
coherent_response = max(responses, key=lambda x: len(x.split()))
message_history.append(coherent_response)
sorted_responses = sorted(responses, key=lambda x: len(x.split()))
# Seleccionar una respuesta aleatoria entre las mejor calificadas
best_responses = sorted_responses[-min(3, len(sorted_responses)):] # Tomar las 3 mejores respuestas o menos
response_text = random.choice(best_responses)
# Guardar el mensaje del usuario y la respuesta en el historial
message_history.append(q)
message_history.append(response_text)
return response_text
# Función para ejecutar la aplicación sin reiniciarla
def run_app():
uvicorn.run(app, host='0.0.0.0', port=7860)
# Ejecutar la aplicación
if __name__ == "__main__":
run_app()