Update app.py

app.py

@@ -147,8 +147,25 @@ class S3DirectStream:
         logger.info(f"Tokenizer loaded for {model_prefix}.")
         return tokenizer
 
-    def download_and_upload_to_s3(self, model_prefix, revision="main"):
+    def download_and_upload_to_s3(self, model_prefix): # Removed revision parameter
         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"
+
         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")