Ggggggg

Paused

App Files Files Community

Hjgugugjhuhjggg commited on Feb 6

Commit

c9f727d

verified ·

1 Parent(s): a542983

Update app.py

Browse files

Files changed (1) hide show

app.py +8 -35

app.py CHANGED Viewed

@@ -27,7 +27,6 @@ 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 = ""
@@ -42,7 +41,7 @@ class GenerateRequest(BaseModel):
     do_sample: bool = True
     stream: bool = True  # Streaming por defecto
     chunk_token_limit: int = 2  # Máximo 2 tokens por bloque
-    token_timeout: float = 0.0  # Timeout en 0: sin timeout
     stop_sequences: list[str] = []
     include_duckasgo: bool = False
@@ -73,8 +72,7 @@ async def load_global_model():
 @app.get("/", response_class=HTMLResponse)
 async def index():
     """
-    Endpoint raíz que devuelve una página HTML simple para permitir la navegación
-    mientras se generan respuestas en paralelo.
     """
     html_content = """
     <html>
@@ -83,7 +81,6 @@ async def index():
         </head>
         <body>
             <h1>Bienvenido al Generador de Texto</h1>
-            <p>El sistema utiliza streaming por defecto para generar respuestas rápidamente.</p>
             <p>Prueba los endpoints <code>/generate</code> o <code>/duckasgo</code>.</p>
         </body>
     </html>
@@ -100,7 +97,6 @@ async def health():
 async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
     """
     Realiza una búsqueda en DuckDuckGo y retorna un resumen de los resultados.
-    Se ejecuta en cada llamada sin almacenar resultados en caché.
     """
     try:
         with DDGS() as ddgs:
@@ -121,8 +117,6 @@ async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
 def generate_next_token(input_ids, past_key_values, gen_config, device):
     """
     Función síncrona que genera el siguiente token utilizando el modelo.
-    Retorna además el log-probability del token seleccionado.
-    Se invoca en paralelo mediante asyncio.to_thread.
     """
     with torch.no_grad():
         outputs = global_model(
@@ -152,15 +146,14 @@ def generate_next_token(input_ids, past_key_values, gen_config, device):
 async def stream_text(request: GenerateRequest, device: str):
     """
     Genera texto de forma streaming, enviando cada bloque con hasta 'chunk_token_limit' tokens.
-    Se continúa en un loop hasta detectar el token de finalización (eos) o alcanzar un límite total.
     """
     global global_model, global_tokenizer, global_tokens
-    # Prepara la entrada y configura la generación
     encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
     input_ids = encoded_input.input_ids
     gen_config = GenerationConfig(
         temperature=request.temperature,
-        max_new_tokens=request.max_new_tokens,  # Límite total opcional
         top_p=request.top_p,
         top_k=request.top_k,
         repetition_penalty=request.repetition_penalty,
@@ -176,44 +169,28 @@ async def stream_text(request: GenerateRequest, device: str):
     max_total_tokens = request.max_new_tokens if request.max_new_tokens > 0 else 1000
     while True:
-        if request.token_timeout > 0:
-            try:
-                next_token, past_key_values, token_logprob = await asyncio.wait_for(
-                    asyncio.to_thread(generate_next_token, input_ids, past_key_values, gen_config, device),
-                    timeout=request.token_timeout
-                )
-            except asyncio.TimeoutError:
-                yield "data: " + json.dumps({"generated_text": "[Token generation timeout, continuing...]"} ) + "\n\n"
-                continue
-        else:
-            next_token, past_key_values, token_logprob = 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)
         current_chunk += token_text
         chunk_token_count += 1
         token_count += 1
-        # Envía el bloque actual cuando se alcanza el límite de tokens por bloque
         if chunk_token_count >= request.chunk_token_limit:
             yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
             current_chunk = ""
             chunk_token_count = 0
-        await asyncio.sleep(0)
         input_ids = next_token
-        # Condición de terminación: se detecta el token eos o se alcanza el límite total
         if token_id == global_tokens["eos_token_id"] or token_count >= max_total_tokens:
             break
-    # Envía cualquier bloque parcial pendiente
     if current_chunk:
         yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
-    # Incluye resultados de búsqueda si se solicita
     if request.include_duckasgo:
         search_summary = await perform_duckasgo_search(request.input_text)
         yield "data: " + json.dumps({"generated_text": search_summary}) + "\n\n"
@@ -222,7 +199,6 @@ async def stream_text(request: GenerateRequest, device: str):
 def synchronous_generation(encoded_input, gen_config, device):
     """
     Función síncrona para generación completa en modo no streaming.
-    Se ejecuta en paralelo mediante asyncio.to_thread.
     """
     with torch.no_grad():
         output = global_model.generate(
@@ -238,7 +214,7 @@ def synchronous_generation(encoded_input, gen_config, device):
 async def generate_text(request: GenerateRequest, background_tasks: BackgroundTasks):
     """
     Endpoint para la generación de texto.
-    En modo streaming se envían bloques de hasta 'chunk_token_limit' tokens hasta finalizar.
     En modo no streaming se devuelve la respuesta completa dividida en bloques.
     """
     global global_model, global_tokenizer, global_tokens
@@ -273,7 +249,6 @@ async def generate_text(request: GenerateRequest, background_tasks: BackgroundTa
                         full_generated_text = full_generated_text.split(stop_seq)[0]
                         break
-            # Dividir la respuesta en bloques de tokens según chunk_token_limit
             final_token_ids = global_tokenizer.encode(full_generated_text, add_special_tokens=False)
             chunks = []
             for i in range(0, len(final_token_ids), request.chunk_token_limit):
@@ -295,7 +270,6 @@ async def generate_text(request: GenerateRequest, background_tasks: BackgroundTa
 async def duckasgo_search(request: DuckasgoRequest):
     """
     Endpoint para búsquedas en DuckDuckGo.
-    Se ejecuta en cada petición sin almacenar resultados en caché.
     """
     try:
         with DDGS() as ddgs:
@@ -310,7 +284,6 @@ def run_server():
     uvicorn.run(app, host="0.0.0.0", port=7860)
 if __name__ == "__main__":
-    # Inicia el servidor en un hilo separado para permitir tareas concurrentes.
     server_thread = threading.Thread(target=run_server, daemon=True)
     server_thread.start()
     while True:

     gc.collect()
     if device == "cuda":
         torch.cuda.empty_cache()
 class GenerateRequest(BaseModel):
     input_text: str = ""
     do_sample: bool = True
     stream: bool = True  # Streaming por defecto
     chunk_token_limit: int = 2  # Máximo 2 tokens por bloque
+    # Se eliminan token_timeout y demás esperas
     stop_sequences: list[str] = []
     include_duckasgo: bool = False
 @app.get("/", response_class=HTMLResponse)
 async def index():
     """
+    Endpoint raíz que devuelve una página HTML simple.
     """
     html_content = """
     <html>
         </head>
         <body>
             <h1>Bienvenido al Generador de Texto</h1>
             <p>Prueba los endpoints <code>/generate</code> o <code>/duckasgo</code>.</p>
         </body>
     </html>
 async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
     """
     Realiza una búsqueda en DuckDuckGo y retorna un resumen de los resultados.
     """
     try:
         with DDGS() as ddgs:
 def generate_next_token(input_ids, past_key_values, gen_config, device):
     """
     Función síncrona que genera el siguiente token utilizando el modelo.
     """
     with torch.no_grad():
         outputs = global_model(
 async def stream_text(request: GenerateRequest, device: str):
     """
     Genera texto de forma streaming, enviando cada bloque con hasta 'chunk_token_limit' tokens.
+    Se continúa generando hasta detectar el token de finalización (eos) o alcanzar un límite total.
     """
     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
     gen_config = GenerationConfig(
         temperature=request.temperature,
+        max_new_tokens=request.max_new_tokens,
         top_p=request.top_p,
         top_k=request.top_k,
         repetition_penalty=request.repetition_penalty,
     max_total_tokens = request.max_new_tokens if request.max_new_tokens > 0 else 1000
     while True:
+        next_token, past_key_values, token_logprob = 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)
         current_chunk += token_text
         chunk_token_count += 1
         token_count += 1
         if chunk_token_count >= request.chunk_token_limit:
             yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
             current_chunk = ""
             chunk_token_count = 0
         input_ids = next_token
         if token_id == global_tokens["eos_token_id"] or token_count >= max_total_tokens:
             break
     if current_chunk:
         yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
     if request.include_duckasgo:
         search_summary = await perform_duckasgo_search(request.input_text)
         yield "data: " + json.dumps({"generated_text": search_summary}) + "\n\n"
 def synchronous_generation(encoded_input, gen_config, device):
     """
     Función síncrona para generación completa en modo no streaming.
     """
     with torch.no_grad():
         output = global_model.generate(
 async def generate_text(request: GenerateRequest, background_tasks: BackgroundTasks):
     """
     Endpoint para la generación de texto.
+    En modo streaming se envían bloques de hasta 'chunk_token_limit' tokens.
     En modo no streaming se devuelve la respuesta completa dividida en bloques.
     """
     global global_model, global_tokenizer, global_tokens
                         full_generated_text = full_generated_text.split(stop_seq)[0]
                         break
             final_token_ids = global_tokenizer.encode(full_generated_text, add_special_tokens=False)
             chunks = []
             for i in range(0, len(final_token_ids), request.chunk_token_limit):
 async def duckasgo_search(request: DuckasgoRequest):
     """
     Endpoint para búsquedas en DuckDuckGo.
     """
     try:
         with DDGS() as ddgs:
     uvicorn.run(app, host="0.0.0.0", port=7860)
 if __name__ == "__main__":
     server_thread = threading.Thread(target=run_server, daemon=True)
     server_thread.start()
     while True: