Update app.py

app.py

@@ -9,6 +9,8 @@ from transformers import (
     AutoModelForCausalLM,
     AutoTokenizer,
     GenerationConfig,
+    StoppingCriteria,
+    StoppingCriteriaList
 )
 import boto3
 import uvicorn
@@ -43,7 +45,7 @@ class GenerateRequest(BaseModel):
     input_text: str = ""
     task_type: str
     temperature: float = 1.0
-    max_new_tokens: int = 3
+    max_new_tokens: int = 3   
     stream: bool = True
     top_p: float = 1.0
     top_k: int = 50    
@@ -92,33 +94,44 @@ class S3ModelLoader:
             )
                 
             tokenizer = AutoTokenizer.from_pretrained(
-                s3_uri, config=config, local_files_only=False
+                s3_uri, config=config, local_files_only=False, padding_side="left"
             )
+            
+            eos_token_id = tokenizer.eos_token_id
+            pad_token_id = tokenizer.pad_token_id
+            eos_token = tokenizer.eos_token
+            pad_token = tokenizer.pad_token
+            padding = tokenizer.padding_side
+                
+            if eos_token_id is not None and pad_token_id is None:
+                pad_token_id = config.pad_token_id or eos_token_id
+                tokenizer.pad_token_id = pad_token_id
 
-            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
-            model_cache[model_name] = (model, tokenizer)
-            return model, tokenizer
+            model_cache[model_name] = (model, tokenizer,eos_token_id,
+            pad_token_id,eos_token,pad_token,padding)
+            return model, tokenizer,eos_token_id,pad_token_id,eos_token,pad_token,padding
        except (EnvironmentError, NoCredentialsError):
             try:
                 config = AutoConfig.from_pretrained(
                     model_name, token=HUGGINGFACE_HUB_TOKEN
                 )
                 tokenizer = AutoTokenizer.from_pretrained(
-                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
+                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN, padding_side="left"
                 )
                 
                 model = AutoModelForCausalLM.from_pretrained(
                     model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
+                
+                eos_token_id = tokenizer.eos_token_id
+                pad_token_id = tokenizer.pad_token_id
+                eos_token = tokenizer.eos_token
+                pad_token = tokenizer.pad_token
+                padding = tokenizer.padding_side
 
-
-                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
+                if eos_token_id is not None and pad_token_id is None:
+                    pad_token_id = config.pad_token_id or eos_token_id
+                    tokenizer.pad_token_id = pad_token_id
 
                 
                 model.save_pretrained(s3_uri)
@@ -134,10 +147,22 @@ class S3ModelLoader:
                 )
                 
                 tokenizer = AutoTokenizer.from_pretrained(
-                    s3_uri, config=config, local_files_only=False
+                    s3_uri, config=config, local_files_only=False, padding_side="left"
                 )
-                model_cache[model_name] = (model, tokenizer)
-                return model, tokenizer
+                
+                eos_token_id = tokenizer.eos_token_id
+                pad_token_id = tokenizer.pad_token_id
+                eos_token = tokenizer.eos_token
+                pad_token = tokenizer.pad_token
+                padding = tokenizer.padding_side
+
+                if eos_token_id is not None and pad_token_id is None:
+                    pad_token_id = config.pad_token_id or eos_token_id
+                    tokenizer.pad_token_id = pad_token_id
+                
+                model_cache[model_name] = (model, tokenizer,eos_token_id,
+                pad_token_id,eos_token,pad_token,padding)
+                return model, tokenizer,eos_token_id,pad_token_id,eos_token,pad_token,padding
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Error loading model: {e}"
@@ -145,6 +170,31 @@ class S3ModelLoader:
 
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 
+class StopOnSequencesCriteria(StoppingCriteria):
+    def __init__(self, stop_sequences, tokenizer):
+        self.stop_sequences = stop_sequences
+        self.tokenizer = tokenizer
+    
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
+        
+        decoded_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
+
+        for seq in self.stop_sequences:
+            if seq in decoded_text:
+                return True
+        return False
+    
+async def generate_stream(model, tokenizer, input_text,
+                             generation_config, stop_sequences,
+                             device, pad_token_id, max_model_length,
+                             max_new_tokens):
+    async def stream():
+        async for token in stream_text(model, tokenizer, input_text,
+                             generation_config, stop_sequences,
+                             device,pad_token_id, max_model_length, max_new_tokens):
+            yield token
+    return stream()
+
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
@@ -161,7 +211,7 @@ async def generate(request: GenerateRequest):
         do_sample = request.do_sample
         stop_sequences = request.stop_sequences
 
-        model, tokenizer = await model_loader.load_model_and_tokenizer(model_name)
+        model, tokenizer, eos_token_id, pad_token_id, eos_token, pad_token, padding = await model_loader.load_model_and_tokenizer(model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
         
@@ -174,14 +224,30 @@ async def generate(request: GenerateRequest):
                 repetition_penalty=repetition_penalty,
                 do_sample=do_sample,
                 num_return_sequences=num_return_sequences,
+                pad_token_id=pad_token_id if pad_token_id is not None else None
             )
+            
+            max_model_length = model.config.max_position_embeddings
+            input_text = input_text[:max_model_length]
 
-            return StreamingResponse(
-                stream_text(model, tokenizer, input_text,
+            streams = [
+                 generate_stream(model, tokenizer, input_text,
                              generation_config, stop_sequences,
-                             device),
+                             device,pad_token_id, max_model_length, max_new_tokens)
+             for _ in range(num_return_sequences)
+            ]
+
+            
+            async def stream_response():
+                async for results in asyncio.as_completed(streams):
+                    async for chunk in await results:
+                        yield chunk
+
+            return StreamingResponse(
+               stream_response(),
                 media_type="text/plain"
             )
+            
         else:
             return HTTPException(status_code=400, detail="Task type not text-to-text")
 
@@ -193,11 +259,13 @@ async def generate(request: GenerateRequest):
 
 async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
-                        device):
-
+                        device,pad_token_id, max_model_length, max_new_tokens):
+    
     encoded_input = tokenizer(
         input_text, return_tensors="pt",
-        truncation=True
+        truncation=True,
+        padding = "max_length",
+        max_length=max_model_length
     ).to(device)
     
     stop_regex = re.compile(r'[\.\?\!\n]+')
@@ -216,6 +284,8 @@ async def stream_text(model, tokenizer, input_text,
     
     
     output_text = ""
+    stop_criteria = StoppingCriteriaList([StopOnSequencesCriteria(stop_sequences, tokenizer)])
+
     while True:
         outputs = model.generate(
             **encoded_input,
@@ -228,8 +298,10 @@ async def stream_text(model, tokenizer, input_text,
             num_return_sequences=generation_config.num_return_sequences,
             output_scores=True,
             return_dict_in_generate=True,
+            pad_token_id=pad_token_id if pad_token_id is not None else None,
+            stopping_criteria = stop_criteria
         )
-
+        
         new_text = tokenizer.decode(
             outputs.sequences[0][len(encoded_input["input_ids"][0]):],
             skip_special_tokens=True
@@ -243,22 +315,32 @@ async def stream_text(model, tokenizer, input_text,
             final_output = output_text[:stop_index]
             
             for text in final_output.split():
-                 yield json.dumps({"text": text, "is_end": False}) + "\n"
+                yield json.dumps({"text": text, "is_end": False}) + "\n"
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
         else:
-            for text in new_text.split():
-              yield json.dumps({"text": text, "is_end": False}) + "\n"
+           
+            tokens = new_text.split()
+            
+            
+            for i in range(0, len(tokens), max_new_tokens):
+                chunk = tokens[i:i + max_new_tokens]
+                chunk_text = " ".join(chunk)
+                for text in chunk_text.split():
+                    yield json.dumps({"text": text, "is_end": False}) + "\n"
         
         if len(new_text) == 0:
+           
             for text in output_text.split():
-                 yield json.dumps({"text": text, "is_end": False}) + "\n"
+                yield json.dumps({"text": text, "is_end": False}) + "\n"
             yield json.dumps({"text": "", "is_end": True}) + "\n"
             break
-
+            
         encoded_input = tokenizer(
-            output_text, return_tensors="pt",
-            truncation=True
+           output_text, return_tensors="pt",
+            truncation=True,
+            padding = "max_length" ,
+            max_length = max_model_length
         ).to(device)
         output_text = ""