Update app.py

app.py

@@ -1,8 +1,8 @@
 from fastapi import FastAPI
+from langchain_community.llms import VLLM
+from langchain_community.cache import GPTCache
 import torch
 from langchain.chains.llm import LLMChain
-from langchain.llms import VLLM
-from langchain.cache import GPTCache
 from transformers import pipeline
 import uvicorn
 import threading
@@ -14,10 +14,8 @@ import psutil
 import os
 import gc
 import logging
-from tqdm import tqdm
 
 logging.basicConfig(level=logging.INFO)
-
 nltk.download('punkt')
 nltk.download('stopwords')
 
@@ -29,24 +27,22 @@ else:
     device = torch.device("cpu")
 
 modelos = {
-    "gpt2-medium": VLLM(model="gpt2-medium"),
-    "qwen2.5-0.5b": VLLM(model="Qwen/Qwen2.5-0.5B-Instruct"),
-    "llamaxd": VLLM(model="Hjgugugjhuhjggg/llama-3.2-1B-spinquant-hf")
+    "gpt2-medium": VLLM(model="gpt2-medium", device=device),
+    "qwen2.5-0.5b": VLLM(model="Qwen/Qwen2.5-0.5B-Instruct", device=device),
+    "llamaxd": VLLM(model="Hjgugugjhuhjggg/llama-3.2-1B-spinquant-hf", device=device)
 }
 
-for nombre, modelo in tqdm(modelos.items()):
-    modelos[nombre] = modelo(to=device)
-
 caches = {
-    nombre: GPTCache(modelo, max_size=1000) for nombre, modelo in modelos.items()
+    nombre: GPTCache(modelo, max_size=1000) 
+    for nombre, modelo in modelos.items()
 }
 
 cadenas = {
-    nombre: LLMChain(modelo, caché) for nombre, modelo, caché in zip(modelos.keys(), modelos.values(), caches.values())
+    nombre: LLMChain(modelo, cache) 
+    for nombre, modelo, cache in zip(modelos.keys(), modelos.values(), caches.values())
 }
 
 summarizer = pipeline("summarization", device=device)
-
 vectorizer = TfidfVectorizer()
 
 def keep_alive():
@@ -56,7 +52,6 @@ def keep_alive():
                 cadena.ask("¿Cuál es el sentido de la vida?")
             except Exception as e:
                 logging.error(f"Error en modelo {cadena}: {e}")
-                cadenas.pop(cadena)
         time.sleep(300)
 
 def liberar_recursos():
@@ -70,7 +65,7 @@ def liberar_recursos():
                     os.kill(proc.info['pid'], 9)
         time.sleep(60)
 
-# Crear hilos para mantener la aplicación viva y liberar recursos
+# Iniciar hilos en segundo plano
 threading.Thread(target=keep_alive, daemon=True).start()
 threading.Thread(target=liberar_recursos, daemon=True).start()
 
@@ -82,7 +77,7 @@ async def pregunta(pregunta: str, modelo: str):
             mensajes = []
             palabras = respuesta.split()
             mensaje_actual = ""
-            for palabra in tqdm(palabras):
+            for palabra in palabras:
                 if len(mensaje_actual.split()) + len(palabra.split()) > 2048:
                     mensajes.append(mensaje_actual)
                     mensaje_actual = palabra
@@ -96,8 +91,8 @@ async def pregunta(pregunta: str, modelo: str):
             respuesta_vec = vectorizer.transform([respuesta])
             similitud = cosine_similarity(pregunta_vec, respuesta_vec)
             return {
-                "respuesta": respuesta,
-                "resumen": resumen[0]["summary_text"],
+                "respuesta": respuesta, 
+                "resumen": resumen[0]["summary_text"], 
                 "similitud": similitud[0][0]
             }
     except Exception as e: