Update app.py

app.py

@@ -6,7 +6,7 @@ from pydantic import BaseModel, field_validator
 from transformers import (
     AutoConfig,
     pipeline,
-    AutoModelForSeq2SeqLM, # Changed AutoModelForCausalLM to AutoModelForSeq2SeqLM
+    AutoModelForSeq2SeqLM,
     AutoTokenizer,
     GenerationConfig,
     StoppingCriteriaList
@@ -69,7 +69,7 @@ class S3ModelLoader:
         s3_uri = self._get_s3_uri(model_name)
         try:
             config = AutoConfig.from_pretrained(s3_uri, local_files_only=True)
-            model = AutoModelForSeq2SeqLM.from_pretrained(s3_uri, config=config, local_files_only=True) # Changed AutoModelForCausalLM
+            model = AutoModelForSeq2SeqLM.from_pretrained(s3_uri, config=config, local_files_only=True)
             tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config, local_files_only=True)
 
             if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
@@ -80,7 +80,7 @@ class S3ModelLoader:
             try:
                 config = AutoConfig.from_pretrained(model_name)
                 tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
-                model = AutoModelForSeq2SeqLM.from_pretrained(model_name, config=config) # Changed AutoModelForCausalLM
+                model = AutoModelForSeq2SeqLM.from_pretrained(model_name, config=config)
 
                 if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
                     tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
@@ -135,6 +135,7 @@ async def generate(request: GenerateRequest):
         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):
@@ -159,38 +160,58 @@ async def stream_text(model, tokenizer, input_text,
                 return last_index + len(seq)
 
         return -1
-
     
     output_text = ""
     
     while True:
         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,
-        output_scores=True,
-        return_dict_in_generate=True,
+            **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,
+            output_scores=True,
+            return_dict_in_generate=True,
         )
     
         new_text = tokenizer.decode(outputs.sequences[0][len(encoded_input["input_ids"][0]):], skip_special_tokens=True)
         
         output_text += new_text
         
-        yield new_text
-        await asyncio.sleep(chunk_delay)
-            
+        
         
         stop_index = find_stop(output_text, stop_sequences)
+        
         if stop_index != -1:
-            yield output_text[:stop_index]
+            final_output = output_text[:stop_index]
+            
+            
+            
+            chunked_output = [final_output[i:i+10] for i in range(0, len(final_output), 10)]
+            
+            for chunk in chunked_output:
+                yield chunk
+                await asyncio.sleep(chunk_delay)
+                
             break
+                
+        else:
+            chunked_output = [new_text[i:i+10] for i in range(0, len(new_text), 10)]
+            for chunk in chunked_output:
+                yield chunk
+                await asyncio.sleep(chunk_delay)
+                
         
         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 chunk in chunked_output:
+                yield chunk
+                await asyncio.sleep(chunk_delay)
             break
 
         encoded_input = tokenizer(output_text,