Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse, JSONResponse
+from fastapi.responses import StreamingResponse
 from pydantic import BaseModel, field_validator
 from transformers import (
     AutoConfig,
@@ -10,7 +10,7 @@ from transformers import (
     AutoTokenizer,
     GenerationConfig,
     StoppingCriteria,
-    StoppingCriteriaList,
+    StoppingCriteriaList
 )
 import boto3
 import uvicorn
@@ -20,6 +20,7 @@ import json
 from huggingface_hub import login
 import base64
 from botocore.exceptions import NoCredentialsError
+import re
 
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
@@ -44,10 +45,10 @@ class GenerateRequest(BaseModel):
     input_text: str = ""
     task_type: str
     temperature: float = 1.0
-    max_new_tokens: int = 3
-    stream: bool = True  # Set default stream to True
+    max_new_tokens: int = 3   
+    stream: bool = True
     top_p: float = 1.0
-    top_k: int = 50
+    top_k: int = 50    
     repetition_penalty: float = 1.0
     num_return_sequences: int = 1
     do_sample: bool = True
@@ -93,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)
@@ -135,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}"
@@ -146,6 +170,37 @@ 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():
+         past_key_values = None
+         input_ids = None
+         async for token,past_key_values_response,input_ids_response, is_end in stream_text(model, tokenizer, input_text,
+                             generation_config, stop_sequences,
+                             device,pad_token_id, max_model_length, max_new_tokens, past_key_values, input_ids):
+            past_key_values = past_key_values_response
+            input_ids = input_ids_response
+            if is_end:
+                break
+            yield token
+    return stream()
+
 @app.post("/generate")
 async def generate(request: GenerateRequest):
     try:
@@ -154,7 +209,7 @@ async def generate(request: GenerateRequest):
         task_type = request.task_type
         temperature = request.temperature
         max_new_tokens = request.max_new_tokens
-        stream = request.stream
+        stream = True
         top_p = request.top_p
         top_k = request.top_k
         repetition_penalty = request.repetition_penalty
@@ -162,10 +217,10 @@ 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)
-
+        
         if "text-to-text" == task_type:
             generation_config = GenerationConfig(
                 temperature=temperature,
@@ -175,20 +230,29 @@ async def generate(request: GenerateRequest):
                 repetition_penalty=repetition_penalty,
                 do_sample=do_sample,
                 num_return_sequences=num_return_sequences,
-                eos_token_id = tokenizer.eos_token_id
+                pad_token_id=pad_token_id if pad_token_id is not None else None
+            )
+            
+            max_model_length = 3
+            input_text = input_text[:max_model_length]
+
+            streams = [
+                 generate_stream(model, tokenizer, input_text,
+                             generation_config, stop_sequences,
+                             device,pad_token_id, max_model_length, max_new_tokens)
+             for _ in range(num_return_sequences)
+            ]
+
+            
+            async def stream_response():
+                for stream in asyncio.as_completed(streams):
+                      async for chunk in await stream:
+                        yield chunk
+            return StreamingResponse(
+               stream_response(),
+                media_type="text/plain"
             )
-            if stream:
-                return StreamingResponse(
-                        stream_text(model, tokenizer, input_text,
-                                     generation_config, stop_sequences,
-                                     device),
-                        media_type="text/plain"
-                    )
-            else:
-                result = await generate_text(model, tokenizer, input_text,
-                                     generation_config, stop_sequences,
-                                     device)
-                return JSONResponse({"text": result, "is_end": True})
+            
         else:
             return HTTPException(status_code=400, detail="Task type not text-to-text")
 
@@ -197,110 +261,112 @@ async def generate(request: GenerateRequest):
             status_code=500, detail=f"Internal server error: {str(e)}"
         )
 
-class StopOnSequences(StoppingCriteria):
-    def __init__(self, stop_sequences, tokenizer):
-        self.stop_sequences = stop_sequences
-        self.tokenizer = tokenizer
-        self.stop_ids = [tokenizer.encode(seq, add_special_tokens=False) for seq in stop_sequences]
-
-    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 stop_sequence in self.stop_sequences:
-             if stop_sequence in decoded_text:
-                 return True
-        return False
 
 async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
-                        device):
+                        device,pad_token_id, max_model_length, max_new_tokens,
+                        past_key_values, input_ids):
     
-    encoded_input = tokenizer(
-        input_text, return_tensors="pt",
-        truncation=True
-    ).to(device)
     
-    stop_criteria = StopOnSequences(stop_sequences, tokenizer)
-    stopping_criteria = StoppingCriteriaList([stop_criteria])
-
-    output_text = ""
+    stop_regex = re.compile(r'[\.\?\!\n]+')
     
-    while True:
+    def find_stop(output_text, stop_sequences):
+        for seq in stop_sequences:
+           if seq in output_text:
+              last_index = output_text.rfind(seq)
+              return last_index + len(seq)
         
-        outputs = await asyncio.to_thread(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,
-                                            stopping_criteria=stopping_criteria
-                                         )
-    
-        new_text = tokenizer.decode(
-            outputs.sequences[0][len(encoded_input["input_ids"][0]):],
-            skip_special_tokens=True
-        )
-
-        if len(new_text) == 0:
-            if not stop_criteria(outputs.sequences, None):
-                for text in output_text.split():
-                    yield json.dumps({"text": text, "is_end": False}) + "\n"
-                yield json.dumps({"text": "", "is_end": True}) + "\n"
-            break
-        
-        output_text += new_text
-        
-        for text in new_text.split():
-              yield json.dumps({"text": text, "is_end": False}) + "\n"
-              
-        if stop_criteria(outputs.sequences, None):
-            yield json.dumps({"text": "", "is_end": True}) + "\n"
-            break
+        match = stop_regex.search(output_text)
+        if match:
+           return match.end()
         
-        encoded_input = tokenizer(
-            output_text, return_tensors="pt",
-            truncation=True
-        ).to(device)
-        output_text = ""
-        
-
-async def generate_text(model, tokenizer, input_text,
-                        generation_config, stop_sequences,
-                        device):
-    encoded_input = tokenizer(
-        input_text, return_tensors="pt",
-        truncation=True
-    ).to(device)
+        return -1
     
-    stop_criteria = StopOnSequences(stop_sequences, tokenizer)
-    stopping_criteria = StoppingCriteriaList([stop_criteria])
-    
-    outputs = await asyncio.to_thread(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,
-                                        stopping_criteria=stopping_criteria
-                                     )
     
+    output_text = ""
+    stop_criteria = StoppingCriteriaList([StopOnSequencesCriteria(stop_sequences, tokenizer)])
     
-    generated_text = tokenizer.decode(
-        outputs.sequences[0], skip_special_tokens=True
-    )
+    if input_ids is None:
+      encoded_input = tokenizer(
+                input_text, return_tensors="pt",
+                truncation=True,
+                padding = "max_length",
+                 max_length=max_model_length
+            ).to(device)
+      input_ids = encoded_input.input_ids
+    else:
+      encoded_input = {
+       "input_ids":input_ids,
+       "past_key_values": past_key_values
+      }
     
-    return generated_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,
+          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
+      )
+      
+      output_text += new_text
+      
+      stop_index = find_stop(output_text, stop_sequences)
+      
+      is_end = False
+      if stop_index != -1 or (hasattr(outputs, "sequences") and outputs.sequences[0][-1] == tokenizer.eos_token_id):
+          final_output = output_text[:stop_index] if stop_index != -1 else output_text
+          
+          for text in final_output.split():
+            yield json.dumps({"text": text, "is_end": False, "temperature": generation_config.temperature, "top_p": generation_config.top_p, "top_k": generation_config.top_k}) + "\n", \
+                    outputs.past_key_values if hasattr(outputs, "past_key_values") else None , \
+                    outputs.sequences if hasattr(outputs, "sequences") else None, True
+          
+          yield json.dumps({"text": "", "is_end": True, "temperature": generation_config.temperature, "top_p": generation_config.top_p, "top_k": generation_config.top_k}) + "\n",\
+                outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
+                  outputs.sequences if hasattr(outputs, "sequences") else None, True
+          break
+      else:
+          
+        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, "temperature": generation_config.temperature, "top_p": generation_config.top_p, "top_k": generation_config.top_k}) + "\n", \
+                    outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
+                      outputs.sequences if hasattr(outputs, "sequences") else None, False
+      
+      if len(new_text) == 0:
+        
+          for text in output_text.split():
+               yield json.dumps({"text": text, "is_end": False, "temperature": generation_config.temperature, "top_p": generation_config.top_p, "top_k": generation_config.top_k}) + "\n", \
+                      outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
+                      outputs.sequences if hasattr(outputs, "sequences") else None, True
+          yield json.dumps({"text": "", "is_end": True, "temperature": generation_config.temperature, "top_p": generation_config.top_p, "top_k": generation_config.top_k}) + "\n",\
+                  outputs.past_key_values if hasattr(outputs, "past_key_values") else None, \
+                    outputs.sequences if hasattr(outputs, "sequences") else None, True
+          break
+      
+      past_key_values = outputs.past_key_values if hasattr(outputs, "past_key_values") else None
+      input_ids = outputs.sequences if hasattr(outputs, "sequences") else None
+        
+      output_text = ""
+
 
 
 @app.post("/generate-image")