Ggggggg

Paused

App Files Files Community

Hjgugugjhuhjggg commited on Feb 6

Commit

fe94180

verified ·

1 Parent(s): 3016722

Update app.py

Browse files

Files changed (1) hide show

app.py +98 -36

app.py CHANGED Viewed

@@ -4,28 +4,34 @@ import torch
 import asyncio
 import threading
 import time
-from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import uvicorn
 from duckduckgo_search import DDGS
 MODEL_NAME = "lilmeaty/my_xdd"
 global_model = None
 global_tokenizer = None
 global_tokens = {}
-# Diccionarios para cachear respuestas
 global_response_cache = {}
 duckasgo_response_cache = {}
 async def cleanup_memory(device: str):
     gc.collect()
     if device == "cuda":
         torch.cuda.empty_cache()
     await asyncio.sleep(0.01)
 class GenerateRequest(BaseModel):
     input_text: str = ""
     max_new_tokens: int = 200
@@ -39,6 +45,7 @@ class GenerateRequest(BaseModel):
     stop_sequences: list[str] = []
     include_duckasgo: bool = False
 class DuckasgoRequest(BaseModel):
     query: str
@@ -46,6 +53,9 @@ app = FastAPI()
 @app.on_event("startup")
 async def load_global_model():
     global global_model, global_tokenizer, global_tokens
     config = AutoConfig.from_pretrained(MODEL_NAME)
     global_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, config=config)
@@ -60,8 +70,28 @@ async def load_global_model():
     global_tokens["pad_token_id"] = global_tokenizer.pad_token_id
     print(f"Modelo {MODEL_NAME} cargado correctamente en {device}.")
 async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
-    # Primero se revisa la caché
     if query in duckasgo_response_cache:
         return duckasgo_response_cache[query]
     try:
@@ -78,11 +108,38 @@ async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
             url = res.get("href", "Sin URL")
             snippet = res.get("body", "")
             result_text += f"{idx}. {title}\n   URL: {url}\n   {snippet}\n"
-    # Guardar en caché
     duckasgo_response_cache[query] = result_text
     return result_text
 async def stream_text(request: GenerateRequest, device: str):
     global global_model, global_tokenizer, global_tokens
     encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
     input_ids = encoded_input.input_ids
@@ -99,24 +156,10 @@ async def stream_text(request: GenerateRequest, device: str):
     )
     past_key_values = None
     for _ in range(request.max_new_tokens):
-        with torch.no_grad():
-            outputs = global_model(
-                input_ids,
-                past_key_values=past_key_values,
-                use_cache=True,
-                return_dict=True
-            )
-        logits = outputs.logits[:, -1, :]
-        past_key_values = outputs.past_key_values
-        if gen_config.do_sample:
-            logits = logits / gen_config.temperature
-            if gen_config.top_k and gen_config.top_k > 0:
-                topk_values, _ = torch.topk(logits, k=gen_config.top_k)
-                logits[logits < topk_values[:, [-1]]] = -float('Inf')
-            probs = torch.nn.functional.softmax(logits, dim=-1)
-            next_token = torch.multinomial(probs, num_samples=1)
-        else:
-            next_token = torch.argmax(logits, dim=-1, keepdim=True)
         token_id = next_token.item()
         token_text = global_tokenizer.decode([token_id], skip_special_tokens=True)
         accumulated_text += token_text
@@ -143,8 +186,27 @@ async def stream_text(request: GenerateRequest, device: str):
         yield "\n" + search_summary
     await cleanup_memory(device)
 @app.post("/generate")
-async def generate_text(request: GenerateRequest):
     global global_model, global_tokenizer, global_tokens, global_response_cache
     if global_model is None or global_tokenizer is None:
         raise HTTPException(status_code=500, detail="El modelo no se ha cargado correctamente.")
@@ -157,9 +219,7 @@ async def generate_text(request: GenerateRequest):
         repetition_penalty=request.repetition_penalty,
         do_sample=request.do_sample,
     )
-    # Solo se cachea si no es un request de stream
     if not request.stream:
-        # Se construye una clave con los parámetros relevantes
         cache_key = (
             request.input_text,
             request.max_new_tokens,
@@ -175,18 +235,15 @@ async def generate_text(request: GenerateRequest):
             return {"generated_text": global_response_cache[cache_key]}
     try:
         if request.stream:
             generator = stream_text(request, device)
             return StreamingResponse(generator, media_type="text/plain")
         else:
             encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
-            with torch.no_grad():
-                output = global_model.generate(
-                    **encoded_input,
-                    generation_config=gen_config,
-                    return_dict_in_generate=True,
-                    output_scores=True,
-                    return_legacy_cache=True  # Se agrega el parámetro solicitado
-                )
             input_length = encoded_input["input_ids"].shape[-1]
             generated_text = global_tokenizer.decode(
                 output.sequences[0][input_length:],
@@ -201,14 +258,17 @@ async def generate_text(request: GenerateRequest):
                 search_summary = await perform_duckasgo_search(request.input_text)
                 generated_text += "\n" + search_summary
             await cleanup_memory(device)
-            # Se almacena en caché la respuesta generada
             global_response_cache[cache_key] = generated_text
             return {"generated_text": generated_text}
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error durante la generación: {e}")
 @app.post("/duckasgo")
 async def duckasgo_search(request: DuckasgoRequest):
     global duckasgo_response_cache
     if request.query in duckasgo_response_cache:
         return JSONResponse(content={"query": request.query, "results": duckasgo_response_cache[request.query]})
@@ -217,13 +277,15 @@ async def duckasgo_search(request: DuckasgoRequest):
             results = ddgs.text(request.query, max_results=10)
         if not results:
             results = []
-        # Se almacena la respuesta en caché
         duckasgo_response_cache[request.query] = results
         return JSONResponse(content={"query": request.query, "results": results})
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error en la búsqueda: {e}")
 def run_server():
     uvicorn.run(app, host="0.0.0.0", port=7860)
 if __name__ == "__main__":

 import asyncio
 import threading
 import time
+from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
 from pydantic import BaseModel
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import uvicorn
 from duckduckgo_search import DDGS
+from concurrent.futures import ThreadPoolExecutor
+# Nombre del modelo a cargar y variables globales
 MODEL_NAME = "lilmeaty/my_xdd"
 global_model = None
 global_tokenizer = None
 global_tokens = {}
+# Caché para respuestas y búsquedas
 global_response_cache = {}
 duckasgo_response_cache = {}
+# Executor para ejecutar tareas en paralelo
+executor = ThreadPoolExecutor(max_workers=4)
 async def cleanup_memory(device: str):
     gc.collect()
     if device == "cuda":
         torch.cuda.empty_cache()
     await asyncio.sleep(0.01)
+# Modelo de datos para generación
 class GenerateRequest(BaseModel):
     input_text: str = ""
     max_new_tokens: int = 200
     stop_sequences: list[str] = []
     include_duckasgo: bool = False
+# Modelo de datos para búsqueda DuckDuckGo
 class DuckasgoRequest(BaseModel):
     query: str
 @app.on_event("startup")
 async def load_global_model():
+    """
+    Carga el modelo y tokenizador en el inicio de la aplicación.
+    """
     global global_model, global_tokenizer, global_tokens
     config = AutoConfig.from_pretrained(MODEL_NAME)
     global_tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, config=config)
     global_tokens["pad_token_id"] = global_tokenizer.pad_token_id
     print(f"Modelo {MODEL_NAME} cargado correctamente en {device}.")
+@app.get("/", response_class=HTMLResponse)
+async def index():
+    """
+    Endpoint raíz que permite visitar el sitio mientras se ejecutan otras operaciones.
+    """
+    html_content = """
+    <html>
+        <head>
+            <title>Mi Sitio de Generación de Texto</title>
+        </head>
+        <body>
+            <h1>Bienvenido a Mi Sitio</h1>
+            <p>Puedes enviar peticiones a <code>/generate</code> o <code>/duckasgo</code> sin afectar la navegación.</p>
+        </body>
+    </html>
+    """
+    return HTMLResponse(content=html_content, status_code=200)
 async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
+    """
+    Realiza búsqueda en DuckDuckGo y utiliza caché para acelerar consultas repetidas.
+    """
     if query in duckasgo_response_cache:
         return duckasgo_response_cache[query]
     try:
             url = res.get("href", "Sin URL")
             snippet = res.get("body", "")
             result_text += f"{idx}. {title}\n   URL: {url}\n   {snippet}\n"
     duckasgo_response_cache[query] = result_text
     return result_text
+def generate_next_token(input_ids, past_key_values, gen_config, device):
+    """
+    Función síncrona que genera el siguiente token usando el modelo.
+    Se invoca en paralelo mediante asyncio.to_thread.
+    """
+    with torch.no_grad():
+        outputs = global_model(
+            input_ids,
+            past_key_values=past_key_values,
+            use_cache=True,
+            return_dict=True
+        )
+    logits = outputs.logits[:, -1, :]
+    past_key_values = outputs.past_key_values
+    if gen_config.do_sample:
+        logits = logits / gen_config.temperature
+        if gen_config.top_k and gen_config.top_k > 0:
+            topk_values, _ = torch.topk(logits, k=gen_config.top_k)
+            logits[logits < topk_values[:, [-1]]] = -float('Inf')
+        probs = torch.nn.functional.softmax(logits, dim=-1)
+        next_token = torch.multinomial(probs, num_samples=1)
+    else:
+        next_token = torch.argmax(logits, dim=-1, keepdim=True)
+    return next_token, past_key_values
 async def stream_text(request: GenerateRequest, device: str):
+    """
+    Genera texto de forma streaming, utilizando generación paralela para cada token.
+    """
     global global_model, global_tokenizer, global_tokens
     encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
     input_ids = encoded_input.input_ids
     )
     past_key_values = None
     for _ in range(request.max_new_tokens):
+        # Ejecuta la generación del siguiente token en paralelo
+        next_token, past_key_values = await asyncio.to_thread(
+            generate_next_token, input_ids, past_key_values, gen_config, device
+        )
         token_id = next_token.item()
         token_text = global_tokenizer.decode([token_id], skip_special_tokens=True)
         accumulated_text += token_text
         yield "\n" + search_summary
     await cleanup_memory(device)
+def synchronous_generation(encoded_input, gen_config, device, request: GenerateRequest):
+    """
+    Función síncrona para la generación completa en modo no streaming.
+    Se ejecuta en paralelo mediante asyncio.to_thread.
+    """
+    with torch.no_grad():
+        output = global_model.generate(
+            **encoded_input,
+            generation_config=gen_config,
+            return_dict_in_generate=True,
+            output_scores=True,
+            return_legacy_cache=True  # Parámetro agregado para optimizar el uso de caché interno
+        )
+    return output
 @app.post("/generate")
+async def generate_text(request: GenerateRequest, background_tasks: BackgroundTasks):
+    """
+    Endpoint para la generación de texto. Se utiliza paralelismo para la generación (streaming o completa)
+    y se cachean respuestas para acelerar peticiones repetidas.
+    """
     global global_model, global_tokenizer, global_tokens, global_response_cache
     if global_model is None or global_tokenizer is None:
         raise HTTPException(status_code=500, detail="El modelo no se ha cargado correctamente.")
         repetition_penalty=request.repetition_penalty,
         do_sample=request.do_sample,
     )
     if not request.stream:
         cache_key = (
             request.input_text,
             request.max_new_tokens,
             return {"generated_text": global_response_cache[cache_key]}
     try:
         if request.stream:
+            # Generación en modo streaming en paralelo
             generator = stream_text(request, device)
             return StreamingResponse(generator, media_type="text/plain")
         else:
             encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
+            # Se ejecuta la generación completa en un hilo paralelo
+            output = await asyncio.to_thread(
+                synchronous_generation, encoded_input, gen_config, device, request
+            )
             input_length = encoded_input["input_ids"].shape[-1]
             generated_text = global_tokenizer.decode(
                 output.sequences[0][input_length:],
                 search_summary = await perform_duckasgo_search(request.input_text)
                 generated_text += "\n" + search_summary
             await cleanup_memory(device)
             global_response_cache[cache_key] = generated_text
+            background_tasks.add_task(lambda: print("Generación completada y cacheada."))
             return {"generated_text": generated_text}
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error durante la generación: {e}")
 @app.post("/duckasgo")
 async def duckasgo_search(request: DuckasgoRequest):
+    """
+    Endpoint para búsquedas en DuckDuckGo con cacheo para acelerar respuestas.
+    """
     global duckasgo_response_cache
     if request.query in duckasgo_response_cache:
         return JSONResponse(content={"query": request.query, "results": duckasgo_response_cache[request.query]})
             results = ddgs.text(request.query, max_results=10)
         if not results:
             results = []
         duckasgo_response_cache[request.query] = results
         return JSONResponse(content={"query": request.query, "results": results})
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error en la búsqueda: {e}")
 def run_server():
+    """
+    Inicia el servidor Uvicorn.
+    """
     uvicorn.run(app, host="0.0.0.0", port=7860)
 if __name__ == "__main__":