Update app.py

app.py

@@ -1,13 +1,14 @@
 import os
 import gc
 import json
+import random
 import torch
 import asyncio
 import threading
 import time
 from fastapi import FastAPI, HTTPException, BackgroundTasks
 from fastapi.responses import StreamingResponse, JSONResponse, HTMLResponse
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig
 import uvicorn
 from duckduckgo_search import DDGS
@@ -30,14 +31,17 @@ async def cleanup_memory(device: str):
 
 class GenerateRequest(BaseModel):
     input_text: str = ""
-    max_new_tokens: int = 2
-    temperature: float = 1.0
-    top_p: float = 1.0
-    top_k: int = 50
-    repetition_penalty: float = 1.0
+    max_new_tokens: int = 200
+    temperature: float = Field(default_factory=lambda: round(random.uniform(0.5, 0.8), 2))
+    top_p: float = Field(default_factory=lambda: round(random.uniform(0.75, 0.95), 2))
+    top_k: int = Field(default_factory=lambda: random.randint(20, 60))
+    repetition_penalty: float = Field(default_factory=lambda: round(random.uniform(1.1, 1.8), 2))
+    frequency_penalty: float = Field(default_factory=lambda: round(random.uniform(0.2, 0.7), 2))
+    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
+    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
@@ -148,24 +152,27 @@ 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 2 tokens generados de forma independiente.
-    Cada respuesta se envía inmediatamente con el campo "generated_text" que contiene solo esos tokens (sin acumulación).
-    Si token_timeout es mayor que 0 se aplica un timeout; de lo contrario se espera sin límite.
+    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
-    # Se crea un objeto de configuración para la generación
+    # Se configura el GenerationConfig con los nuevos parámetros aleatorios
     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,
     )
-    # Variable para acumulación interna (para validación de stop sequences)
-    all_tokens = ""
-    # Variable para el bloque actual (máximo 2 tokens por bloque)
+    # 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
@@ -185,17 +192,10 @@ async def stream_text(request: GenerateRequest, device: str):
             )
         token_id = next_token.item()
         token_text = global_tokenizer.decode([token_id], skip_special_tokens=True)
-        all_tokens += token_text
+        # En lugar de acumular, se envía únicamente el token generado
         current_chunk += token_text
         chunk_token_count += 1
-        # Validar si se encuentra alguna secuencia de parada en el texto acumulado (para detener la generación)
-        if request.stop_sequences:
-            for stop_seq in request.stop_sequences:
-                if stop_seq in all_tokens:
-                    yield "data: " + json.dumps({"generated_text": current_chunk}) + "\n\n"
-                    await cleanup_memory(device)
-                    return
-        # Si se han generado 2 tokens en el bloque, se envía el bloque actual y se reinicia
+        # Si se han generado 2 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 = ""
@@ -204,7 +204,7 @@ async def stream_text(request: GenerateRequest, device: str):
         input_ids = next_token
         if token_id == global_tokens["eos_token_id"]:
             break
-    # Si quedó un bloque parcial, se envía
+    # 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
@@ -239,14 +239,19 @@ async def generate_text(request: GenerateRequest, background_tasks: BackgroundTa
     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 (se usarán los parámetros ya asignados en 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)
     try:
         if request.stream:
             generator = stream_text(request, device)