File size: 4,320 Bytes
c9c8569 c6ab5e0 c9c8569 c6ab5e0 c9c8569 c6ab5e0 c9c8569 c6ab5e0 c9c8569 c6ab5e0 c9c8569 d228ca8 c9c8569 c6ab5e0 c9c8569 c6ab5e0 c9c8569 f7ba4cc c9c8569 f7ba4cc c9c8569 f7ba4cc c9c8569 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
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 transformers import pipeline
import uvicorn
import threading
import time
import nltk
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity
import psutil
import os
import gc
import logging
logging.basicConfig(level=logging.INFO)
nltk.download('punkt')
nltk.download('stopwords')
app = FastAPI()
if torch.cuda.is_available():
device = torch.device("cuda")
else:
device = torch.device("cpu")
modelos = {
"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)
}
caches = {
nombre: GPTCache(modelo, max_size=1000)
for nombre, modelo in modelos.items()
}
cadenas = {
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():
while True:
for cadena in cadenas.values():
try:
cadena.ask("¿Cuál es el sentido de la vida?")
except Exception as e:
logging.error(f"Error en modelo {cadena}: {e}")
time.sleep(300)
def liberar_recursos():
while True:
memoria_ram = psutil.virtual_memory().available / (1024.0 ** 3)
espacio_disco = psutil.disk_usage('/').free / (1024.0 ** 3)
if memoria_ram < 5 or espacio_disco < 5:
gc.collect()
for proc in psutil.process_iter(['pid', 'name']):
if proc.info['name'] == 'python':
os.kill(proc.info['pid'], 9)
time.sleep(60)
# Iniciar hilos en segundo plano
threading.Thread(target=keep_alive, daemon=True).start()
threading.Thread(target=liberar_recursos, daemon=True).start()
@app.post("/pregunta")
async def pregunta(pregunta: str, modelo: str):
try:
respuesta = cadenas[modelo].ask(pregunta)
if len(respuesta.split()) > 2048:
mensajes = []
palabras = respuesta.split()
mensaje_actual = ""
for palabra in palabras:
if len(mensaje_actual.split()) + len(palabra.split()) > 2048:
mensajes.append(mensaje_actual)
mensaje_actual = palabra
else:
mensaje_actual += " " + palabra
mensajes.append(mensaje_actual)
return {"respuestas": mensajes}
else:
resumen = summarizer(respuesta, max_length=50, min_length=5, do_sample=False)
pregunta_vec = vectorizer.fit_transform([pregunta])
respuesta_vec = vectorizer.transform([respuesta])
similitud = cosine_similarity(pregunta_vec, respuesta_vec)
return {
"respuesta": respuesta,
"resumen": resumen[0]["summary_text"],
"similitud": similitud[0][0]
}
except Exception as e:
logging.error(f"Error en modelo {modelo}: {e}")
return {"error": f"Modelo {modelo} no disponible"}
@app.post("/resumen")
async def resumen(texto: str):
try:
resumen = summarizer(texto, max_length=50, min_length=5, do_sample=False)
return {"resumen": resumen[0]["summary_text"]}
except Exception as e:
logging.error(f"Error en resumen: {e}")
return {"error": "Error en resumen"}
@app.post("/similitud")
async def similitud(texto1: str, texto2: str):
try:
texto1_vec = vectorizer.fit_transform([texto1])
texto2_vec = vectorizer.transform([texto2])
similitud = cosine_similarity(texto1_vec, texto2_vec)
return {"similitud": similitud[0][0]}
except Exception as e:
logging.error(f"Error en similitud: {e}")
return {"error": "Error en similitud"}
@app.get("/modelos")
async def modelos():
return {"modelos": list(cadenas.keys())}
@app.get("/estado")
async def estado():
return {"estado": "activo"}
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=8000)
|