Update app.py

app.py

@@ -38,7 +38,6 @@ class DownloadModelRequest(BaseModel):
     model_id: str
     pipeline_task: str
     input_text: str
-    revision: str = "main"
 
     @field_validator('model_id')
     def validate_model_id(cls, value):
@@ -75,9 +74,14 @@ class S3DirectStream:
             logger.info(f"File {key} does not exist in S3.")
             return False
 
-    def load_model_from_stream(self, model_prefix, revision):
+    def load_model_from_stream(self, model_prefix):
         try:
-            logger.info(f"Loading model {model_prefix} (revision {revision})...")
+            logger.info(f"Loading model {model_prefix}...")
+            revision = self._get_latest_revision(model_prefix)
+            if revision is None:
+                logger.error(f"Could not determine revision for {model_prefix}")
+                raise ValueError(f"Could not determine revision for {model_prefix}")
+
             if self.file_exists_in_s3(f"{model_prefix}/config.json") and \
                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...")
@@ -86,7 +90,7 @@ class S3DirectStream:
             logger.info(f"Model {model_prefix} not found in S3. Downloading and uploading...")
             self.download_and_upload_to_s3(model_prefix, revision)
             logger.info(f"Downloaded and uploaded {model_prefix}. Loading from S3...")
-            return self.load_model_from_stream(model_prefix, revision)
+            return self.load_model_from_stream(model_prefix)
         except HTTPException as e:
             logger.error(f"Error loading model: {e}")
             return None
@@ -165,12 +169,15 @@ class S3DirectStream:
 
         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()
+            logger.info(f"Hugging Face API Response: Status Code = {index_response.status_code}, Headers = {index_response.headers}")
             index_content = index_response.text
+            logger.info(f"Index content: {index_content}")
             model_files = [f for f in index_content.split('\n') if f.endswith(('.bin', '.safetensors'))]
             return model_files
         except requests.exceptions.RequestException as e:
@@ -198,19 +205,27 @@ class S3DirectStream:
                 logger.error(f"Error downloading from {url}: Status code {response.status_code}")
                 raise HTTPException(status_code=500, detail=f"Error downloading file from {url}")
 
+    def _get_latest_revision(self, model_prefix):
+        try:
+            api = HfApi()
+            model_info = api.model_info(model_prefix)
+            return model_info.default_revision
+        except Exception as e:
+            logger.error(f"Error getting latest revision for {model_prefix}: {e}")
+            return None
+
 
 @app.post("/predict/")
 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
-        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)
+        model = streamer.load_model_from_stream(model_id)
 
         if model is None:
             logger.error(f"Failed to load model {model_id}")