Update app.py

app.py

@@ -5,7 +5,6 @@ from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, field_validator
 from transformers import (
     AutoConfig,
-    pipeline,
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
@@ -130,15 +129,8 @@ async def generate(request: GenerateRequest):
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
-async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay, max_length=2048):
-    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
-    input_length = encoded_input["input_ids"].shape[1]
-    remaining_tokens = max_length - input_length
-
-    if remaining_tokens <= 0:
-        yield ""
-
-    generation_config.max_new_tokens = min(remaining_tokens, generation_config.max_new_tokens)
+async def stream_text(model, tokenizer, input_text, generation_config, stop_sequences, device, chunk_delay):
+    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True).to(device)
 
     def stop_criteria(input_ids, scores):
         decoded_output = tokenizer.decode(int(input_ids[0][-1]), skip_special_tokens=True)
@@ -146,7 +138,7 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
 
     stopping_criteria = StoppingCriteriaList([stop_criteria])
 
-    output_text = ""
+    token_buffer = []
     outputs = model.generate(
         **encoded_input,
         do_sample=generation_config.do_sample,
@@ -158,19 +150,29 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
         num_return_sequences=generation_config.num_return_sequences,
         stopping_criteria=stopping_criteria,
         output_scores=True,
-        return_dict_in_generate=True
+        return_dict_in_generate=True,
+        streamer=None  # Ensure streamer is None for manual token processing
     )
 
     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)  # Simula el delay entre tokens
+            token_buffer.append(token)
+            if len(token_buffer) >= 10:
+                yield "".join(token_buffer)
+                token_buffer = []
+            await asyncio.sleep(chunk_delay)
+
+        if token_buffer:
+            yield "".join(token_buffer)
+            token_buffer = []
 
-        if stop_sequences and any(stop in output_text for stop in stop_sequences):
-            yield output_text
+        if stop_sequences and any(stop in tokenizer.decode(output, skip_special_tokens=True) for stop in stop_sequences):
             return
 
+        encoded_input = tokenizer.build_inputs_with_special_tokens(output)
+        encoded_input = {'input_ids': torch.tensor([encoded_input]).to(device)}
+
         outputs = model.generate(
             **encoded_input,
             do_sample=generation_config.do_sample,
@@ -182,7 +184,8 @@ async def stream_text(model, tokenizer, input_text, generation_config, stop_sequ
             num_return_sequences=generation_config.num_return_sequences,
             stopping_criteria=stopping_criteria,
             output_scores=True,
-            return_dict_in_generate=True
+            return_dict_in_generate=True,
+            streamer=None
         )
 
 @app.post("/generate-image")
@@ -190,7 +193,7 @@ async def generate_image(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        
+
         image_generator = pipeline("text-to-image", model=validated_body.model_name, device=device)
         image = image_generator(validated_body.input_text)[0]
 
@@ -208,7 +211,7 @@ async def generate_text_to_speech(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        
+
         audio_generator = pipeline("text-to-speech", model=validated_body.model_name, device=device)
         audio = audio_generator(validated_body.input_text)[0]