Update app.py

app.py

@@ -147,39 +147,27 @@ class S3DirectStream:
         logger.info(f"Tokenizer loaded for {model_prefix}.")
         return tokenizer
 
-    def download_and_upload_to_s3(self, model_prefix): # Removed revision parameter
+    def download_and_upload_to_s3(self, model_prefix):
         logger.info(f"Downloading and uploading model files for {model_prefix} to S3...")
         try:
             api = HfApi()
-            model_info = api.model_info(model_prefix)
-            if hasattr(model_info, 'revision'):
-                revision = model_info.revision
-                if revision:
-                    revision = revision
-                else:
-                    logger.warning(f"No revision found for {model_prefix}, using 'main'")
-                    revision = "main"
-            else:
-                logger.warning(f"ModelInfo object for {model_prefix} does not have a 'revision' attribute, using 'main'")
-                revision = "main"
-        except Exception as e:
-            logger.error(f"Error getting latest revision for {model_prefix}: {e}")
-            revision = "main"
+            model_files = api.list_repo_files(model_prefix)
 
-        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")
+            for file_info in model_files:
+                if file_info.rfilename.endswith(('.bin', '.safetensors', 'config.json', 'tokenizer.json')):
+                    file_url = api.download_file(model_prefix, file_info.rfilename)
+                    s3_key = f"{model_prefix}/{file_info.rfilename}"
+                    self.download_and_upload_to_s3_url(file_url, s3_key)
+                    logger.info(f"Downloaded and uploaded {s3_key}")
 
-        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}")
+            logger.info(f"Finished downloading and uploading model files for {model_prefix}.")
 
-        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")
-
-        logger.info(f"Finished downloading and uploading model files for {model_prefix}.")
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Error downloading model files from HuggingFace: {e}")
+            raise HTTPException(status_code=500, detail=f"Error downloading model files from Hugging Face") from e
+        except Exception as e:
+            logger.error(f"An unexpected error occurred: {e}")
+            raise HTTPException(status_code=500, detail=f"An unexpected error occurred during model download") from e
 
 
     def _get_model_files(self, model_prefix, revision="main"):