Update app.py

app.py

@@ -114,30 +114,31 @@ class GCSModelLoader:
                 return AutoConfig.from_pretrained(pretrained_model_name_or_path="", _commit_hash=None, config_dict=json.loads(config_content), trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
             except Exception as e:
                 logger.error(f"Error loading config from GCS: {e}")
-        else:
+        try:
             logger.info(f"Downloading config from Hugging Face for {model_name}")
-            try:
-                config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
-                gcs_model_folder = self._get_gcs_uri(model_name)
-                self._create_model_folder(model_name)
-                self._upload_content(json.dumps(config.to_dict()).encode('utf-8'), f"{gcs_model_folder}/config.json")
-                return config
-            except Exception as e:
-                logger.error(f"Error loading config from Hugging Face: {e}")
-                return None
+            config = AutoConfig.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
+            gcs_model_folder = self._get_gcs_uri(model_name)
+            self._create_model_folder(model_name)
+            self._upload_content(json.dumps(config.to_dict()).encode('utf-8'), f"{gcs_model_folder}/config.json")
+            return config
+        except Exception as e:
+            logger.error(f"Error loading config from Hugging Face: {e}")
+            return None
 
     def load_tokenizer(self, model_name):
         gcs_tokenizer_path = self._get_gcs_uri(model_name)
         tokenizer_files = ["tokenizer_config.json", "vocab.json", "merges.txt", "tokenizer.json", "special_tokens_map.json"]
-        if all(self._blob_exists(f"{gcs_tokenizer_path}/{f}") for f in tokenizer_files):
+        gcs_files_exist = all(self._blob_exists(f"{gcs_tokenizer_path}/{f}") for f in tokenizer_files)
+
+        if gcs_files_exist:
             try:
                 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:
-            logger.info(f"Downloading tokenizer from Hugging Face for {model_name}")
             try:
+                logger.info(f"Downloading tokenizer from Hugging Face for {model_name}")
                 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, token=HUGGINGFACE_HUB_TOKEN)
                 gcs_model_folder = self._get_gcs_uri(model_name)
                 self._create_model_folder(model_name)
@@ -161,7 +162,8 @@ class GCSModelLoader:
                 self._create_model_folder(model_name)
                 for filename in os.listdir(model.config.name_or_path):
                     if filename.endswith((".bin", ".safetensors")):
-                        self._upload_content(open(os.path.join(model.config.name_or_path, filename), 'rb').read(), f"{gcs_model_folder}/{filename}")
+                        blob = self.bucket.blob(f"{gcs_model_folder}/{filename}")
+                        blob.upload_from_filename(os.path.join(model.config.name_or_path, filename))
                 logger.info(f"Model '{model_name}' downloaded from Hugging Face and saved to GCS.")
                 return model
             except Exception as e:
@@ -171,30 +173,36 @@ class GCSModelLoader:
         logger.info(f"Found weight files in GCS for '{model_name}': {weight_files}")
 
         loaded_state_dict = {}
+        error_occurred = False
         for weight_file in weight_files:
             logger.info(f"Streaming weight file from GCS: {weight_file}")
             blob = self.bucket.blob(weight_file)
             try:
-                weight_bytes = blob.download_as_bytes()
+                blob_content = blob.download_as_bytes()
                 if weight_file.endswith(".safetensors"):
-                    loaded_state_dict.update(safe_load(weight_bytes))
+                    loaded_state_dict.update(safe_load(blob_content))
                 else:
-                    loaded_state_dict.update(torch.load(weight_bytes))
+                    loaded_state_dict.update(torch.load(io.BytesIO(blob_content), map_location="cpu"))
             except Exception as e:
                 logger.error(f"Error streaming and loading weights from GCS {weight_file}: {e}")
-                logger.info(f"Attempting to reload model '{model_name}' from Hugging Face due to loading error.")
-                try:
-                    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)
-                    self._create_model_folder(model_name)
-                    for filename in os.listdir(model.config.name_or_path):
-                        if filename.endswith((".bin", ".safetensors")):
-                            self._upload_content(open(os.path.join(model.config.name_or_path, filename), 'rb').read(), f"{gcs_model_folder}/{filename}")
-                    logger.info(f"Model '{model_name}' reloaded from Hugging Face and saved to GCS.")
-                    return model
-                except Exception as redownload_error:
-                    logger.error(f"Error redownloading model from Hugging Face: {redownload_error}")
-                    raise HTTPException(status_code=500, detail=f"Failed to load or redownload model: {redownload_error}")
+                error_occurred = True
+                break
+
+        if error_occurred:
+            logger.info(f"Attempting to reload model '{model_name}' from Hugging Face due to loading error.")
+            try:
+                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)
+                self._create_model_folder(model_name)
+                for filename in os.listdir(model.config.name_or_path):
+                    if filename.endswith((".bin", ".safetensors")):
+                        upload_blob = self.bucket.blob(f"{gcs_model_folder}/{filename}")
+                        upload_blob.upload_from_filename(os.path.join(model.config.name_or_path, filename))
+                logger.info(f"Model '{model_name}' reloaded from Hugging Face and saved to GCS.")
+                return model
+            except Exception as redownload_error:
+                logger.error(f"Error redownloading model from Hugging Face: {redownload_error}")
+                raise HTTPException(status_code=500, detail=f"Failed to load or redownload model: {redownload_error}")
 
         try:
             model = AutoModelForCausalLM.from_config(config, trust_remote_code=True)