Update app.py

app.py

@@ -1,7 +1,7 @@
 import os
 import torch
 from fastapi import FastAPI, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel, field_validator
 from transformers import (
     AutoConfig,
@@ -10,17 +10,15 @@ from transformers import (
     AutoTokenizer,
     GenerationConfig,
     StoppingCriteria,
-    StoppingCriteriaList
+    StoppingCriteriaList,
 )
 import boto3
 import uvicorn
-import asyncio
 from transformers import pipeline
 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")
@@ -45,10 +43,10 @@ class GenerateRequest(BaseModel):
     input_text: str = ""
     task_type: str
     temperature: float = 1.0
-    max_new_tokens: int = 3   
-    stream: bool = True
+    max_new_tokens: int = 3
+    stream: bool = True  # Set default stream to 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
@@ -79,7 +77,7 @@ class S3ModelLoader:
          return f"s3://{self.bucket_name}/" \
                f"{model_name.replace('/', '-')}"
 
-    async def load_model_and_tokenizer(self, model_name):
+    def load_model_and_tokenizer(self, model_name):
        if model_name in model_cache:
           return model_cache[model_name]
        
@@ -94,44 +92,33 @@ class S3ModelLoader:
             )
                 
             tokenizer = AutoTokenizer.from_pretrained(
-                s3_uri, config=config, local_files_only=False, padding_side="left"
+                s3_uri, config=config, local_files_only=False
             )
-            
-            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
+            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
        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, padding_side="left"
+                    model_name, config=config, token=HUGGINGFACE_HUB_TOKEN
                 )
                 
                 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 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.save_pretrained(s3_uri)
@@ -147,22 +134,10 @@ class S3ModelLoader:
                 )
                 
                 tokenizer = AutoTokenizer.from_pretrained(
-                    s3_uri, config=config, local_files_only=False, padding_side="left"
+                    s3_uri, config=config, local_files_only=False
                 )
-                
-                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
+                model_cache[model_name] = (model, tokenizer)
+                return model, tokenizer
             except Exception as e:
                 raise HTTPException(
                     status_code=500, detail=f"Error loading model: {e}"
@@ -170,46 +145,15 @@ 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):
+def generate(request: GenerateRequest):
     try:
         model_name = request.model_name
         input_text = request.input_text
         task_type = request.task_type
         temperature = request.temperature
         max_new_tokens = request.max_new_tokens
-        stream = True
+        stream = request.stream
         top_p = request.top_p
         top_k = request.top_k
         repetition_penalty = request.repetition_penalty
@@ -217,10 +161,10 @@ async def generate(request: GenerateRequest):
         do_sample = request.do_sample
         stop_sequences = request.stop_sequences
 
-        model, tokenizer, eos_token_id, pad_token_id, eos_token, pad_token, padding = await model_loader.load_model_and_tokenizer(model_name)
+        model, tokenizer = 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,
@@ -230,29 +174,20 @@ 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 = 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"
+                eos_token_id = tokenizer.eos_token_id
             )
-            
+            if stream:
+                return StreamingResponse(
+                        stream_text(model, tokenizer, input_text,
+                                     generation_config, stop_sequences,
+                                     device),
+                        media_type="text/plain"
+                    )
+            else:
+                result = 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")
 
@@ -261,116 +196,113 @@ 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]
 
-async def 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):
-    
-    
-    stop_regex = re.compile(r'[\.\?\!\n]+')
-    
-    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)
-        
-        match = stop_regex.search(output_text)
-        if match:
-           return match.end()
+    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
         
-        return -1
+        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
+
+def stream_text(model, tokenizer, input_text,
+                        generation_config, stop_sequences,
+                        device):
     
+    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_criteria = StoppingCriteriaList([StopOnSequencesCriteria(stop_sequences, tokenizer)])
-    
-    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
-      }
     
     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:
+        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,
+                                            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
         
-          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 += 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
+        
+        encoded_input = tokenizer(
+            output_text, return_tensors="pt",
+            truncation=True
+        ).to(device)
+        output_text = ""
         
-      output_text = ""
-
 
+def generate_text(model, tokenizer, input_text,
+                        generation_config, stop_sequences,
+                        device):
+    encoded_input = tokenizer(
+        input_text, return_tensors="pt",
+        truncation=True
+    ).to(device)
+    
+    stop_criteria = StopOnSequences(stop_sequences, tokenizer)
+    stopping_criteria = StoppingCriteriaList([stop_criteria])
+    
+    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,
+                                        stopping_criteria=stopping_criteria
+                                     )
+    
+    
+    generated_text = tokenizer.decode(
+        outputs.sequences[0], skip_special_tokens=True
+    )
+    
+    return generated_text
 
 @app.post("/generate-image")
-async def generate_image(request: GenerateRequest):
+def generate_image(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -398,7 +330,7 @@ async def generate_image(request: GenerateRequest):
 
 
 @app.post("/generate-text-to-speech")
-async def generate_text_to_speech(request: GenerateRequest):
+def generate_text_to_speech(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -429,7 +361,7 @@ async def generate_text_to_speech(request: GenerateRequest):
 
 
 @app.post("/generate-video")
-async def generate_video(request: GenerateRequest):
+def generate_video(request: GenerateRequest):
     try:
         validated_body = request
         device = "cuda" if torch.cuda.is_available() else "cpu"