Update app.py

app.py

@@ -79,7 +79,7 @@ class S3DirectStream:
         try:
             logger.info(f"Loading model {model_prefix} (revision {revision})...")
             if self.file_exists_in_s3(f"{model_prefix}/config.json") and \
-               (self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin") or self.file_exists_in_s3(f"{model_prefix}/model.safetensors")):
+               any(self.file_exists_in_s3(f"{model_prefix}/{file}") for file in self._get_model_files(model_prefix, revision)):
                 logger.info(f"Model {model_prefix} found in S3. Loading...")
                 return self.load_model_from_existing_s3(model_prefix)
 
@@ -98,22 +98,32 @@ class S3DirectStream:
         config = AutoConfig.from_pretrained(model_prefix, **config_dict)
         logger.info(f"Config loaded for {model_prefix}.")
 
-        if self.file_exists_in_s3(f"{model_prefix}/model.safetensors"):
-            logger.info(f"Loading safetensors model for {model_prefix} from S3...")
-            model_stream = self.stream_from_s3(f"{model_prefix}/model.safetensors")
-            model = AutoModelForCausalLM.from_config(config)
-            model.load_state_dict(safetensors.torch.load_stream(model_stream))
-            logger.info(f"Safetensors model loaded for {model_prefix}.")
-        elif self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin"):
-            logger.info(f"Loading PyTorch model for {model_prefix} from S3...")
-            model_stream = self.stream_from_s3(f"{model_prefix}/pytorch_model.bin")
-            model = AutoModelForCausalLM.from_config(config)
-            state_dict = torch.load(model_stream, map_location="cpu")
-            model.load_state_dict(state_dict)
-            logger.info(f"PyTorch model loaded for {model_prefix}.")
-        else:
-            logger.error(f"No model file found for {model_prefix} in S3")
-            raise EnvironmentError(f"No model file found for {model_prefix} in S3")
+        model_files = self._get_model_files(model_prefix)
+        if not model_files:
+            logger.error(f"No model files found for {model_prefix} in S3")
+            raise EnvironmentError(f"No model files found for {model_prefix} in S3")
+
+        state_dict = {}
+        for model_file in model_files:
+            model_path = os.path.join(model_prefix, model_file)
+            logger.info(f"Loading model file: {model_path}")
+            model_stream = self.stream_from_s3(model_path)
+            try:
+                if model_path.endswith(".safetensors"):
+                    shard_state = safetensors.torch.load_stream(model_stream)
+                elif model_path.endswith(".bin"):
+                    shard_state = torch.load(model_stream, map_location="cpu")
+                else:
+                    logger.error(f"Unsupported model file type: {model_path}")
+                    raise ValueError(f"Unsupported model file type: {model_path}")
+
+                state_dict.update(shard_state)
+            except Exception as e:
+                logger.exception(f"Error loading model file {model_path}: {e}")
+                raise
+
+        model = AutoModelForCausalLM.from_config(config)
+        model.load_state_dict(state_dict)
         return model
 
     def load_tokenizer_from_stream(self, model_prefix):
@@ -139,18 +149,37 @@ class S3DirectStream:
         return tokenizer
 
     def download_and_upload_to_s3(self, model_prefix, revision="main"):
-        model_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/pytorch_model.bin"
-        safetensors_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/model.safetensors"
-        tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/tokenizer.json"
+        logger.info(f"Downloading and uploading model files for {model_prefix} to S3...")
         config_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/config.json"
+        self.download_and_upload_to_s3_url(config_url, f"{model_prefix}/config.json")
 
-        logger.info(f"Downloading and uploading model files for {model_prefix} to S3...")
-        self.download_and_upload_to_s3_url(model_url, f"{model_prefix}/pytorch_model.bin")
-        self.download_and_upload_to_s3_url(safetensors_url, f"{model_prefix}/model.safetensors")
+        model_files = self._get_model_files(model_prefix, revision)
+        for model_file in model_files:
+            url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/{model_file}"
+            s3_key = f"{model_prefix}/{model_file}"
+            self.download_and_upload_to_s3_url(url, s3_key)
+            logger.info(f"Downloaded and uploaded {s3_key}")
+
+        tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/tokenizer.json"
         self.download_and_upload_to_s3_url(tokenizer_url, f"{model_prefix}/tokenizer.json")
-        self.download_and_upload_to_s3_url(config_url, f"{model_prefix}/config.json")
+
         logger.info(f"Finished downloading and uploading model files for {model_prefix}.")
 
+    def _get_model_files(self, model_prefix, revision="main"):
+        index_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/"
+        try:
+            index_response = requests.get(index_url)
+            index_response.raise_for_status()
+            index_content = index_response.text
+            model_files = [f for f in index_content.split('\n') if f.endswith(('.bin', '.safetensors'))]
+            return model_files
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Error retrieving model index: {e}")
+            raise HTTPException(status_code=500, detail=f"Error retrieving model files from Hugging Face") from e
+        except (IndexError, ValueError) as e:
+            logger.error(f"Error parsing model file names from Hugging Face: {e}")
+            raise HTTPException(status_code=500, detail=f"Error retrieving model files from Hugging Face") from e
+
     def download_and_upload_to_s3_url(self, url, s3_key):
         logger.info(f"Downloading from {url}...")
         with requests.get(url, stream=True) as response:
@@ -174,24 +203,22 @@ class S3DirectStream:
 async def predict(model_request: DownloadModelRequest):
     try:
         logger.info(f"Received request: Model={model_request.model_id}, Task={model_request.pipeline_task}, Input={model_request.input_text}")
-        model_id = model_request.model_id # Fixed: Use model_id, not model_name
+        model_id = model_request.model_id
         revision = model_request.revision
         task = model_request.pipeline_task
         input_text = model_request.input_text
 
         streamer = S3DirectStream(S3_BUCKET_NAME)
         logger.info("Loading model and tokenizer...")
-        model = streamer.load_model_from_stream(model_id, revision)  # Use model_id
+        model = streamer.load_model_from_stream(model_id, revision)
 
         if model is None:
             logger.error(f"Failed to load model {model_id}")
             raise HTTPException(status_code=500, detail=f"Failed to load model {model_id}")
 
-        tokenizer = streamer.load_tokenizer_from_stream(model_id) # Use model_id
-
+        tokenizer = streamer.load_tokenizer_from_stream(model_id)
         logger.info("Model and tokenizer loaded.")
 
-
         if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering",  "summarization", "zero-shot-classification"]:
             raise HTTPException(status_code=400, detail="Unsupported pipeline task")