Update app.py

app.py

@@ -18,6 +18,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")
@@ -93,13 +94,20 @@ class S3ModelLoader:
             tokenizer = AutoTokenizer.from_pretrained(
                 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
 
-            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(
@@ -112,12 +120,16 @@ class S3ModelLoader:
                 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,8 +147,20 @@ class S3ModelLoader:
                 tokenizer = AutoTokenizer.from_pretrained(
                     s3_uri, config=config, local_files_only=False
                 )
-                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}"
@@ -160,7 +184,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)
         
@@ -178,7 +202,7 @@ async def generate(request: GenerateRequest):
             return StreamingResponse(
                 stream_text(model, tokenizer, input_text,
                              generation_config, stop_sequences,
-                             device),
+                             device,pad_token_id),
                 media_type="text/plain"
             )
         else:
@@ -192,22 +216,28 @@ async def generate(request: GenerateRequest):
 
 async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
-                        device):
+                        device,pad_token_id):
 
     encoded_input = tokenizer(
         input_text, return_tensors="pt",
         truncation=True
     ).to(device)
-
-
+    
+    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)
-
+           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()
+        
         return -1
     
+    
     output_text = ""
     while True:
         outputs = model.generate(
@@ -221,6 +251,7 @@ 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
         )
 
         new_text = tokenizer.decode(
@@ -251,7 +282,8 @@ async def stream_text(model, tokenizer, input_text,
 
         encoded_input = tokenizer(
             output_text, return_tensors="pt",
-            truncation=True
+            truncation=True,
+            padding = "max_length" if pad_token_id is not None else False
         ).to(device)
         output_text = ""