Update app.py

app.py

@@ -96,26 +96,35 @@ class S3DirectStream:
         config = AutoConfig.from_pretrained(model_prefix, **config_dict)
         logger.info(f"Config loaded for {model_prefix}.")
 
-        model_files = self._get_model_files(model_prefix) #This line is not used anymore
-
-        state_dict = {}
-        for model_file in model_files: #This loop is not used anymore
-            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}")
+        try:
+            api = HfApi()
+            model_files = api.list_repo_files(model_prefix)
+            state_dict = {}
+            for file_info in model_files:
+                if file_info.rfilename.endswith(('.bin', '.safetensors')):
+                    file_url = api.download_file(model_prefix, file_info.rfilename)
+                    model_path = os.path.join(model_prefix, file_info.rfilename)
+                    logger.info(f"Downloading model file from {file_url} to {model_path} ...")
+                    with requests.get(file_url, stream=True) as response:
+                        if response.status_code == 200:
+                            try:
+                                model_stream = response.raw
+                                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")
+                                state_dict.update(shard_state)
+                                logger.info(f"Downloaded and loaded model file {model_path}")
+                            except Exception as e:
+                                logger.exception(f"Error loading model file {model_path}: {e}")
+                                raise
+                        else:
+                            logger.error(f"Error downloading {file_url} with status code: {response.status_code}")
+                            raise HTTPException(status_code=500, detail=f"Error downloading model file from Hugging Face")
+        except Exception as e:
+            logger.exception(f"Error loading model files for {model_prefix} : {e}")
+            raise
 
-                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)