Update app.py

app.py

@@ -9,7 +9,6 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
-    StoppingCriteriaList
 )
 import boto3
 import uvicorn
@@ -49,7 +48,6 @@ class GenerateRequest(BaseModel):
     repetition_penalty: float = 1.0
     num_return_sequences: int = 1
     do_sample: bool = True
-    chunk_delay: float = 0.0
     stop_sequences: list[str] = []
 
     @field_validator("model_name")
@@ -139,7 +137,6 @@ async def generate(request: GenerateRequest):
         repetition_penalty = request.repetition_penalty
         num_return_sequences = request.num_return_sequences
         do_sample = request.do_sample
-        chunk_delay = request.chunk_delay
         stop_sequences = request.stop_sequences
 
         model, tokenizer = await model_loader.\
@@ -161,7 +158,7 @@ async def generate(request: GenerateRequest):
             return StreamingResponse(
                 stream_text(model, tokenizer, input_text,
                              generation_config, stop_sequences,
-                             device, chunk_delay),
+                             device),
                 media_type="text/plain"
             )
         else:
@@ -175,7 +172,7 @@ async def generate(request: GenerateRequest):
 
 async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
-                        device, chunk_delay, max_length=2048):
+                        device, max_length=2048):
     encoded_input = tokenizer(
         input_text, return_tensors="pt",
         truncation=True, max_length=max_length
@@ -226,30 +223,22 @@ async def stream_text(model, tokenizer, input_text,
         if stop_index != -1:
             final_output = output_text[:stop_index]
 
-            chunked_output = [final_output[i:i+10]
-                             for i in range(0, len(final_output), 10)]
+            for char in final_output:
+                yield json.dumps({"text": char, "is_end": False}) + "\n"
 
-            for chunk in chunked_output:
-                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
-                await asyncio.sleep(chunk_delay)
 
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
 
         else:
-            chunked_output = [new_text[i:i+10]
-                             for i in range(0, len(new_text), 10)]
-            for chunk in chunked_output:
-                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
-                await asyncio.sleep(chunk_delay)
+
+            for char in new_text:
+                yield json.dumps({"text": char, "is_end": False}) + "\n"
 
         if len(output_text) >= generation_config.max_new_tokens:
-            chunked_output = [output_text[i:i+10]
-                             for i in range(0, len(output_text), 10)]
+            for char in output_text:
+                yield json.dumps({"text": char, "is_end": False}) + "\n"
 
-            for chunk in chunked_output:
-                yield json.dumps({"text": chunk, "is_end": False}) + "\n"
-                await asyncio.sleep(chunk_delay)
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break