Ggggggg

Paused

App Files Files Community

Hjgugugjhuhjggg commited on Feb 6

Commit

28fbc8a

verified ·

1 Parent(s): 5624ca8

Update app.py

Browse files

Files changed (1) hide show

app.py +172 -43

app.py CHANGED Viewed

@@ -7,17 +7,20 @@ import asyncio
 import threading
 import time
 from fastapi import FastAPI, HTTPException, BackgroundTasks
-from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel, Field
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import uvicorn
 from duckduckgo_search import DDGS
 from concurrent.futures import ThreadPoolExecutor
 MODEL_NAME = "lilmeaty/my_xdd"
 global_model = None
 global_tokenizer = None
 global_tokens = {}
 executor = ThreadPoolExecutor(max_workers=4)
 async def cleanup_memory(device: str):
@@ -37,20 +40,28 @@ class GenerateRequest(BaseModel):
     presence_penalty: float = Field(default_factory=lambda: round(random.uniform(0.2, 0.7), 2))
     seed: int = Field(default_factory=lambda: random.randint(0, 1000))
     do_sample: bool = True
-    stream: bool = True
-    chunk_token_limit: int = 2
-    token_timeout: float = 0.0
     stop_sequences: list[str] = []
     include_duckasgo: bool = False
 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)
-    global_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, config=config, torch_dtype=torch.float16)
     if global_tokenizer.eos_token_id is not None and global_tokenizer.pad_token_id is None:
         global_tokenizer.pad_token_id = config.pad_token_id or global_tokenizer.eos_token_id
     device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -59,9 +70,67 @@ async def load_global_model():
     global_tokens["pad_token_id"] = global_tokenizer.pad_token_id
     print(f"Modelo {MODEL_NAME} cargado correctamente en {device}.")
 def generate_next_token(input_ids, past_key_values, gen_config, device):
     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:
@@ -81,27 +150,52 @@ def generate_next_token(input_ids, past_key_values, gen_config, device):
     return next_token, past_key_values, token_logprob.item()
 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
-    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, frequency_penalty=request.frequency_penalty, presence_penalty=request.presence_penalty, do_sample=request.do_sample)
     torch.manual_seed(request.seed)
     current_chunk = ""
     chunk_token_count = 0
     past_key_values = None
     for _ in range(request.max_new_tokens):
         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
         if chunk_token_count >= request.chunk_token_limit:
             yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
             current_chunk = ""
@@ -110,42 +204,54 @@ async def stream_text(request: GenerateRequest, device: str):
         input_ids = next_token
         if token_id == global_tokens["eos_token_id"]:
             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"
     await cleanup_memory(device)
-async def perform_duckasgo_search(query: str, max_results: int = 3) -> str:
-    try:
-        with DDGS() as ddgs:
-            results = ddgs.text(query, max_results=max_results)
-    except Exception as e:
-        return f"Error en la búsqueda de DuckDuckGo: {e}"
-    if not results:
-        result_text = "No se encontraron resultados en DuckDuckGo."
-    else:
-        result_text = "\nResultados de búsqueda (DuckDuckGo):\n"
-        for idx, res in enumerate(results, start=1):
-            title = res.get("title", "Sin título")
-            url = res.get("href", "Sin URL")
-            snippet = res.get("body", "")
-            result_text += f"{idx}. {title}\n   URL: {url}\n   {snippet}\n"
-    return result_text
 def synchronous_generation(encoded_input, gen_config, 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)
     return output
 @app.post("/generate")
 async def generate_text(request: GenerateRequest, background_tasks: BackgroundTasks):
     global global_model, global_tokenizer, global_tokens
     if global_model is None or global_tokenizer is None:
         raise HTTPException(status_code=500, detail="El modelo no se ha cargado correctamente.")
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    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, frequency_penalty=request.frequency_penalty, presence_penalty=request.presence_penalty, do_sample=request.do_sample)
     torch.manual_seed(request.seed)
     try:
         if request.stream:
@@ -155,33 +261,56 @@ async def generate_text(request: GenerateRequest, background_tasks: BackgroundTa
             encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
             output = await asyncio.to_thread(synchronous_generation, encoded_input, gen_config, device)
             input_length = encoded_input["input_ids"].shape[-1]
-            full_text = global_tokenizer.decode(output.sequences[0][input_length:], skip_special_tokens=True)
-            tokens = global_tokenizer.tokenize(full_text)
-            chunks = []
-            for i in range(0, len(tokens), request.chunk_token_limit):
-                chunk_tokens = tokens[i:i+request.chunk_token_limit]
-                chunk_text = global_tokenizer.convert_tokens_to_string(chunk_tokens)
-                chunks.append(chunk_text)
             if request.stop_sequences:
                 for stop_seq in request.stop_sequences:
-                    for idx, chunk in enumerate(chunks):
-                        if stop_seq in chunk:
-                            chunks[idx] = chunk.split(stop_seq)[0]
-                            chunks = chunks[:idx+1]
-                            break
             if request.include_duckasgo:
                 search_summary = await perform_duckasgo_search(request.input_text)
-                chunks.append(search_summary)
             await cleanup_memory(device)
             background_tasks.add_task(lambda: print("Generación completada."))
-            return {"chunks": chunks}
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Error durante la generación: {e}")
 def run_server():
     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:

 import threading
 import time
 from fastapi import FastAPI, HTTPException, BackgroundTasks
+from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
 from pydantic import BaseModel, Field
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import uvicorn
 from duckduckgo_search import DDGS
 from concurrent.futures import ThreadPoolExecutor
+# Nombre del modelo y variables globales
 MODEL_NAME = "lilmeaty/my_xdd"
 global_model = None
 global_tokenizer = None
 global_tokens = {}
+# Executor para ejecutar tareas en paralelo
 executor = ThreadPoolExecutor(max_workers=4)
 async def cleanup_memory(device: str):
     presence_penalty: float = Field(default_factory=lambda: round(random.uniform(0.2, 0.7), 2))
     seed: int = Field(default_factory=lambda: random.randint(0, 1000))
     do_sample: bool = True
+    stream: bool = True  # Streaming por defecto
+    chunk_token_limit: int = 2  # Máximo 2 tokens por respuesta (bloque)
+    token_timeout: float = 0.0  # Timeout en 0: sin timeout
     stop_sequences: list[str] = []
     include_duckasgo: bool = False
+class DuckasgoRequest(BaseModel):
+    query: str
 app = FastAPI()
 @app.on_event("startup")
 async def load_global_model():
+    """
+    Carga el modelo y el tokenizador global al iniciar 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_model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME, config=config, torch_dtype=torch.float16
+    )
     if global_tokenizer.eos_token_id is not None and global_tokenizer.pad_token_id is None:
         global_tokenizer.pad_token_id = config.pad_token_id or global_tokenizer.eos_token_id
     device = "cuda" if torch.cuda.is_available() else "cpu"
     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 devuelve una página HTML simple para permitir la navegación
+    mientras se generan respuestas en paralelo.
+    """
+    html_content = """
+    <html>
+        <head>
+            <title>Generación de Texto - Streaming por defecto</title>
+        </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>
+    """
+    return HTMLResponse(content=html_content, status_code=200)
+@app.get("/health")
+async def health():
+    """
+    Endpoint de salud para verificar el estado del servidor.
+    """
+    return {"status": "ok"}
+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:
+            results = ddgs.text(query, max_results=max_results)
+    except Exception as e:
+        return f"Error en la búsqueda de DuckDuckGo: {e}"
+    if not results:
+        result_text = "No se encontraron resultados en DuckDuckGo."
+    else:
+        result_text = "\nResultados de búsqueda (DuckDuckGo):\n"
+        for idx, res in enumerate(results, start=1):
+            title = res.get("title", "Sin título")
+            url = res.get("href", "Sin URL")
+            snippet = res.get("body", "")
+            result_text += f"{idx}. {title}\n   URL: {url}\n   {snippet}\n"
+    return result_text
 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.
+    Esta función 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:
     return next_token, past_key_values, token_logprob.item()
 async def stream_text(request: GenerateRequest, device: str):
+    """
+    Genera texto de forma streaming, enviando cada bloque con hasta 2 tokens generados de forma independiente.
+    Cada respuesta se envía inmediatamente con el campo "generated_text" que contiene únicamente esos tokens (sin acumulación).
+    Si token_timeout es mayor que 0 se aplica un timeout; de lo contrario, se espera sin límite.
+    """
     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
+    # Configurar GenerationConfig con los parámetros del request
+    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,
+        frequency_penalty=request.frequency_penalty,
+        presence_penalty=request.presence_penalty,
+        do_sample=request.do_sample,
+    )
+    # Fijar la semilla para la generación
     torch.manual_seed(request.seed)
+    # Variables para manejo de bloques de tokens
     current_chunk = ""
     chunk_token_count = 0
     past_key_values = None
     for _ in range(request.max_new_tokens):
         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)
+        # Se acumula el token generado
         current_chunk += token_text
         chunk_token_count += 1
+        # Si se han generado 'chunk_token_limit' tokens en este bloque, se envía el bloque actual y se reinicia
         if chunk_token_count >= request.chunk_token_limit:
             yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
             current_chunk = ""
         input_ids = next_token
         if token_id == global_tokens["eos_token_id"]:
             break
+    # Enviar cualquier bloque parcial pendiente
     if current_chunk:
         yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
+    # Si se solicita incluir búsqueda, se envía al final en el mismo formato
     if request.include_duckasgo:
         search_summary = await perform_duckasgo_search(request.input_text)
         yield "data: " + json.dumps({"generated_text": search_summary}) + "\n\n"
     await cleanup_memory(device)
 def synchronous_generation(encoded_input, gen_config, device):
+    """
+    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
+        )
     return output
 @app.post("/generate")
 async def generate_text(request: GenerateRequest, background_tasks: BackgroundTasks):
+    """
+    Endpoint para la generación de texto.
+    Por defecto se utiliza streaming para emitir cada bloque (máximo 2 tokens) tan pronto estén listos.
+    Cada respuesta incluye únicamente el campo "generated_text" con los tokens generados en ese bloque.
+    En el modo no streaming se divide la respuesta en bloques de tokens, si ésta excede el límite.
+    """
     global global_model, global_tokenizer, global_tokens
     if global_model is None or global_tokenizer is None:
         raise HTTPException(status_code=500, detail="El modelo no se ha cargado correctamente.")
     device = "cuda" if torch.cuda.is_available() else "cpu"
+    # Configuración del GenerationConfig
+    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,
+        frequency_penalty=request.frequency_penalty,
+        presence_penalty=request.presence_penalty,
+        do_sample=request.do_sample,
+    )
+    # Fijar la semilla para la generación
     torch.manual_seed(request.seed)
     try:
         if request.stream:
             encoded_input = global_tokenizer(request.input_text, return_tensors="pt").to(device)
             output = await asyncio.to_thread(synchronous_generation, encoded_input, gen_config, device)
             input_length = encoded_input["input_ids"].shape[-1]
+            # Decodificar la parte generada (excluyendo la entrada)
+            full_generated_text = global_tokenizer.decode(
+                output.sequences[0][input_length:], skip_special_tokens=True
+            )
+            # Aplicar secuencias de corte (stop_sequences) si se han definido
             if request.stop_sequences:
                 for stop_seq in request.stop_sequences:
+                    if stop_seq in full_generated_text:
+                        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):
+                chunk_ids = final_token_ids[i:i+request.chunk_token_limit]
+                chunk_text = global_tokenizer.decode(chunk_ids, skip_special_tokens=True)
+                chunks.append(chunk_text)
+            # Si se solicita incluir búsqueda en DuckDuckGo, se agrega como bloque extra
             if request.include_duckasgo:
                 search_summary = await perform_duckasgo_search(request.input_text)
+                chunks.append("\n" + search_summary)
             await cleanup_memory(device)
             background_tasks.add_task(lambda: print("Generación completada."))
+            return {"generated_chunks": chunks}
     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.
+    Se ejecuta en cada petición sin almacenar la respuesta final en caché.
+    """
+    try:
+        with DDGS() as ddgs:
+            results = ddgs.text(request.query, max_results=10)
+        if not results:
+            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__":
+    # Inicia el servidor en un hilo separado para permitir otras tareas concurrentes.
     server_thread = threading.Thread(target=run_server, daemon=True)
     server_thread.start()
     while True: