Update app.py

app.py

@@ -43,10 +43,10 @@ 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    # Changed back to 50
+    top_k: int = 50    
     repetition_penalty: float = 1.0
     num_return_sequences: int = 1
     do_sample: bool = True
@@ -92,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)
@@ -145,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}"
@@ -184,7 +161,7 @@ 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 = await model_loader.load_model_and_tokenizer(model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
         
@@ -197,17 +174,12 @@ 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,
                              generation_config, stop_sequences,
-                             device,pad_token_id, max_model_length),
+                             device),
                 media_type="text/plain"
             )
         else:
@@ -221,13 +193,11 @@ async def generate(request: GenerateRequest):
 
 async def stream_text(model, tokenizer, input_text,
                         generation_config, stop_sequences,
-                        device,pad_token_id, max_model_length):
+                        device):
 
     encoded_input = tokenizer(
         input_text, return_tensors="pt",
-        truncation=True,
-        padding = "max_length",
-        max_length=max_model_length
+        truncation=True
     ).to(device)
     
     stop_regex = re.compile(r'[\.\?\!\n]+')
@@ -258,7 +228,6 @@ 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
         )
 
         new_text = tokenizer.decode(
@@ -286,16 +255,15 @@ async def stream_text(model, tokenizer, input_text,
                  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,
-            padding = "max_length" ,
-            max_length = max_model_length
+            output_text, return_tensors="pt",
+            truncation=True
         ).to(device)
         output_text = ""
 
 
+
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
     try: