Update app.py

app.py

@@ -74,28 +74,22 @@ class S3ModelLoader:
             logging.info(f"Trying to load {model_name} from S3...")
             config = AutoConfig.from_pretrained(s3_uri)
             model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config)
-            tokenizer = AutoTokenizer.from_pretrained(s3_uri)
+            tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config)
 
             if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
-                if config.pad_token_id is not None:
-                    tokenizer.pad_token_id = config.pad_token_id
-                else:
-                    tokenizer.pad_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
+                tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
 
             logging.info(f"Loaded {model_name} from S3 successfully.")
             return model, tokenizer
         except EnvironmentError:
             logging.info(f"Model {model_name} not found in S3. Downloading...")
             try:
-                model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
-                tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
-                
+                config = AutoConfig.from_pretrained(model_name)
+                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
+                model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN)
+
                 if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
-                    config = AutoConfig.from_pretrained(model_name)
-                    if config.pad_token_id is not None:
-                        tokenizer.pad_token_id = config.pad_token_id
-                    else:
-                        tokenizer.pad_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
+                    tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
 
                 logging.info(f"Downloaded {model_name} successfully.")
                 logging.info(f"Saving {model_name} to S3...")
@@ -128,10 +122,7 @@ async def generate(request: Request, body: GenerateRequest):
                 top_k=validated_body.top_k,
                 repetition_penalty=validated_body.repetition_penalty,
                 do_sample=validated_body.do_sample,
-                num_return_sequences=validated_body.num_return_sequences,
-                stopping_criteria=StoppingCriteriaList(
-                    [lambda _, outputs: tokenizer.decode(outputs[0][-1], skip_special_tokens=True) in validated_body.stop_sequences] if validated_body.stop_sequences else []
-                )
+                num_return_sequences=validated_body.num_return_sequences
             )
 
             async def stream_text():
@@ -149,7 +140,11 @@ async def generate(request: Request, body: GenerateRequest):
 
                     generation_config.max_new_tokens = min(remaining_tokens, validated_body.max_new_tokens)
 
-                    output = model.generate(**encoded_input, generation_config=generation_config)
+                    stopping_criteria = StoppingCriteriaList(
+                        [lambda _, outputs: tokenizer.decode(outputs[0][-1], skip_special_tokens=True) in validated_body.stop_sequences] if validated_body.stop_sequences else []
+                    )
+
+                    output = model.generate(**encoded_input, generation_config=generation_config, stopping_criteria=stopping_criteria)
                     chunk = tokenizer.decode(output[0], skip_special_tokens=True)
                     generated_text += chunk
                     yield chunk