Update app.py

app.py

@@ -73,7 +73,6 @@ class GenerateRequest(BaseModel):
 
     @model_validator(mode='before')
     def set_default_token_ids(cls, values):
-        # These will be populated after tokenizer is loaded, but need defaults to avoid pydantic errors
         values.setdefault("pad_token_id", None)
         values.setdefault("eos_token_id", None)
         values.setdefault("sep_token_id", None)
@@ -89,7 +88,7 @@ class GCSModelLoader:
 
     def _blob_exists(self, blob_path):
         blob = self.bucket.blob(blob_path)
-        return blob.exists(client=self.bucket.client)
+        return blob.exists()
 
     def _download_content(self, blob_path):
         blob = self.bucket.blob(blob_path)
@@ -106,15 +105,15 @@ class GCSModelLoader:
         config_content = self._download_content(gcs_config_path)
         if config_content:
             try:
-                return AutoConfig.from_pretrained(pretrained_model_name_or_path=None, trust_remote_code=True, config_dict=json.loads(config_content))
+                return AutoConfig.from_pretrained(pretrained_model_name_or_path=None, trust_remote_code=True, config_dict=json.loads(config_content), token=HUGGINGFACE_HUB_TOKEN)
             except Exception as e:
                 logger.error(f"Error loading config from GCS: {e}")
                 return None
         else:
             try:
-                config = AutoConfig.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN, trust_remote_code=True)
+                config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
-                self._upload_content(json.dumps(config.to_dict()).encode('utf-8'), f"{gcs_model_folder}/config.json")
+                bucket.blob(f"{gcs_model_folder}/config.json").upload_from_string(json.dumps(config.to_dict()).encode('utf-8'))
                 return config
             except Exception as e:
                 logger.error(f"Error loading config from Hugging Face and saving to GCS: {e}")
@@ -127,18 +126,17 @@ class GCSModelLoader:
 
         if gcs_files_exist:
             try:
-                return AutoTokenizer.from_pretrained(gcs_tokenizer_path, trust_remote_code=True)
+                return AutoTokenizer.from_pretrained(gcs_tokenizer_path, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
             except Exception as e:
                 logger.error(f"Error loading tokenizer from GCS: {e}")
                 return None
         else:
             try:
-                tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN, trust_remote_code=True)
+                tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
-                os.makedirs(gcs_model_folder, exist_ok=True)  # Ensure the folder exists in GCS
-                for file in tokenizer.save_pretrained(gcs_model_folder):
+                for file in tokenizer.save_pretrained(None):
                     with open(file, 'rb') as f:
-                        self._upload_content(f.read(), f"{gcs_model_folder}/{os.path.basename(file)}")
+                        bucket.blob(f"{gcs_model_folder}/{os.path.basename(file)}").upload_from_string(f.read())
                 return tokenizer
             except Exception as e:
                 logger.error(f"Error loading tokenizer from Hugging Face and saving to GCS: {e}")
@@ -157,12 +155,11 @@ class GCSModelLoader:
                 raise HTTPException(status_code=500, detail=f"Error loading model from GCS: {e}")
         else:
             try:
-                model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN, trust_remote_code=True)
+                model = AutoModelForCausalLM.from_pretrained(model_name, config=config, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
-                os.makedirs(gcs_model_folder, exist_ok=True) # Ensure the folder exists in GCS
                 for filename in os.listdir(model.save_pretrained(None)):
-                    with open(os.path.join(model.save_pretrained(None), filename), 'rb') as f:
-                        self._upload_content(f.read(), f"{gcs_model_folder}/{filename}")
+                    with open(os.path.join(model.save_pretrained(None)), 'rb') as f:
+                        bucket.blob(f"{gcs_model_folder}/{filename}").upload_from_string(f.read())
                 return model
             except Exception as e:
                 logger.error(f"Error loading model from Hugging Face and saving to GCS: {e}")
@@ -200,10 +197,9 @@ async def generate(request: GenerateRequest):
 
     try:
         gcs_model_folder_uri = model_loader._get_gcs_uri(model_name)
-        if not model_loader._blob_exists(f"{gcs_model_folder_uri}/config.json"):
+        if not bucket.blob(f"{gcs_model_folder_uri}/config.json").exists():
             logger.info(f"Model '{model_name}' not found in GCS, creating placeholder.")
             bucket.blob(f"{gcs_model_folder_uri}/.placeholder").upload_from_string("")
-
         config = model_loader.load_config(model_name)
         if not config:
             raise HTTPException(status_code=400, detail="Model configuration could not be loaded.")
@@ -212,7 +208,6 @@ async def generate(request: GenerateRequest):
         if not tokenizer:
             raise HTTPException(status_code=400, detail="Tokenizer could not be loaded.")
 
-        # Update token IDs from tokenizer if not provided in request
         if request.pad_token_id is None:
             request.pad_token_id = tokenizer.pad_token_id
         if request.eos_token_id is None:
@@ -233,9 +228,7 @@ async def generate(request: GenerateRequest):
             pad_token_id=request.pad_token_id,
             eos_token_id=request.eos_token_id,
             sep_token_id=request.sep_token_id,
-            unk_token_id=request.unk_token_id,
-            return_dict_in_generate=True,
-            output_scores=True
+            unk_token_id=request.unk_token_id
         )
 
         if task_type == "text-to-text":