File size: 4,300 Bytes
c9c8569 daa26d1 c9c8569 d228ca8 c9c8569 d228ca8 c9c8569 d228ca8 c9c8569 d228ca8 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 133 134 135 136 |
from fastapi import FastAPI
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
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
from tqdm import tqdm
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"),
"qwen2.5-0.5b": VLLM(model="Qwen/Qwen2.5-0.5B-Instruct"),
"llamaxd": VLLM(model="Hjgugugjhuhjggg/llama-3.2-1B-spinquant-hf")
}
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()
}
cadenas = {
nombre: LLMChain(modelo, cachΓ©) for nombre, modelo, cachΓ© 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}")
cadenas.pop(cadena)
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)
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 tqdm(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)
|