Update app.py

app.py

@@ -9,8 +9,7 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
-    StoppingCriteriaList,
-    TextIteratorStreamer # Importar TextIteratorStreamer
+    StoppingCriteriaList
 )
 import uvicorn
 import asyncio
@@ -24,7 +23,7 @@ class GenerateRequest(BaseModel):
     input_text: str = ""
     task_type: str
     temperature: float = 1.0
-    max_new_tokens: int = 200
+    max_new_tokens: int = 4
     stream: bool = True
     top_p: float = 1.0
     top_k: int = 50
@@ -121,27 +120,44 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
 
     stopping_criteria = StoppingCriteriaList([stop_criteria])
 
-    streamer = TextIteratorStreamer(tokenizer, chunk_delay=chunk_delay, skip_prompt=True) # Inicializar streamer
-
-    generation_kwargs = dict(
+    output_text = ""
+    outputs = model.generate(
         **encoded_input,
-        generation_config=generation_config,
+        do_sample=generation_config.do_sample,
+        max_new_tokens=generation_config.max_new_tokens,
+        temperature=generation_config.temperature,
+        top_p=generation_config.top_p,
+        top_k=generation_config.top_k,
+        repetition_penalty=generation_config.repetition_penalty,
+        num_return_sequences=generation_config.num_return_sequences,
         stopping_criteria=stopping_criteria,
-        streamer=streamer, # Pasar streamer a generate
-        return_dict_in_generate=True,
-        output_scores=True
+        output_scores=True,
+        return_dict_in_generate=True
     )
 
-    async def generate_task():
-        model.generate(**generation_kwargs) # Ejecutar generate en background
-
-    asyncio.create_task(generate_task()) # Iniciar la tarea de generación
+    for output in outputs.sequences:
+        for token_id in output:
+            token = tokenizer.decode(token_id, skip_special_tokens=True)
+            yield token
+            await asyncio.sleep(chunk_delay)
 
-    for token in streamer: # Iterar sobre el streamer para obtener tokens uno por uno
-        yield token
-        if stop_sequences and any(stop in token for stop in stop_sequences): # Comprobar stop sequences en cada token
+        if stop_sequences and any(stop in output_text for stop in stop_sequences):
+            yield output_text
             return
 
+        outputs = model.generate(
+            **encoded_input,
+            do_sample=generation_config.do_sample,
+            max_new_tokens=generation_config.max_new_tokens,
+            temperature=generation_config.temperature,
+            top_p=generation_config.top_p,
+            top_k=generation_config.top_k,
+            repetition_penalty=generation_config.repetition_penalty,
+            num_return_sequences=generation_config.num_return_sequences,
+            stopping_criteria=stopping_criteria,
+            output_scores=True,
+            return_dict_in_generate=True
+        )
 
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):