Update app.py

app.py

@@ -1,4 +1,4 @@
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, hf_hub_download
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, field_validator
 import requests
@@ -82,109 +82,104 @@ class S3DirectStream:
                 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...")
-                return self.load_model_from_existing_s3(model_prefix)
+            config = self._load_config(model_prefix, revision)
+            if config is None:
+                logger.error(f"Failed to load config for {model_prefix}")
+                raise ValueError(f"Failed to load config for {model_prefix}")
+
+            model = self._load_model(model_prefix, config, revision)
+            if model is None:
+                logger.error(f"Failed to load model {model_prefix}")
+                raise ValueError(f"Failed to load model {model_prefix}")
+
+            return model
 
-            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)
         except HTTPException as e:
             logger.error(f"Error loading model: {e}")
-            return None
+            raise
+        except Exception as e:
+            logger.exception(f"Unexpected error loading model: {e}")
+            raise HTTPException(status_code=500, detail=f"An unexpected error occurred while loading the model.")
 
-    def load_model_from_existing_s3(self, model_prefix):
-        logger.info(f"Loading config for {model_prefix} from S3...")
-        config_stream = self.stream_from_s3(f"{model_prefix}/config.json")
-        config_dict = json.load(config_stream)
-        config = AutoConfig.from_pretrained(model_prefix, **config_dict)
-        logger.info(f"Config loaded for {model_prefix}.")
-
-        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}")
+    def _load_config(self, model_prefix, revision):
+        try:
+            logger.info(f"Downloading config for {model_prefix} (revision {revision})...")
+            config_path = hf_hub_download(repo_id=model_prefix, filename="config.json", revision=revision)
+            with open(config_path, "r", encoding="utf-8") as f:
+                config_dict = json.load(f)
+            return AutoConfig.from_pretrained(model_prefix, **config_dict)
+        except Exception as e:
+            logger.error(f"Error loading config: {e}")
+            return None
 
-                state_dict.update(shard_state)
-            except Exception as e:
-                logger.exception(f"Error loading model file {model_path}: {e}")
-                raise
+    def _load_model(self, model_prefix, config, revision):
+        try:
+            logger.info(f"Downloading model files for {model_prefix} (revision {revision})...")
+            model_files = self._get_model_files(model_prefix, revision)
+            if not model_files:
+                logger.error(f"No model files found for {model_prefix}")
+                return None
+
+            state_dict = {}
+            for model_file in model_files:
+                logger.info(f"Downloading model file: {model_file}")
+                file_path = hf_hub_download(repo_id=model_prefix, filename=model_file, revision=revision)
+                with open(file_path, "rb") as f:
+                    if model_file.endswith(".safetensors"):
+                        shard_state = safetensors.torch.load_file(file_path)
+                    elif model_file.endswith(".bin"):
+                        shard_state = torch.load(f, map_location="cpu")
+                    else:
+                        logger.error(f"Unsupported model file type: {model_file}")
+                        raise ValueError(f"Unsupported model file type: {model_file}")
+                    state_dict.update(shard_state)
+
+            model = AutoModelForCausalLM.from_config(config)
+            model.load_state_dict(state_dict)
+            return model
 
-        model = AutoModelForCausalLM.from_config(config)
-        model.load_state_dict(state_dict)
-        return model
+        except Exception as e:
+            logger.exception(f"Error loading model: {e}")
+            return None
 
     def load_tokenizer_from_stream(self, model_prefix):
         try:
             logger.info(f"Loading tokenizer for {model_prefix}...")
-            if self.file_exists_in_s3(f"{model_prefix}/tokenizer.json"):
-                logger.info(f"Tokenizer for {model_prefix} found in S3. Loading...")
-                return self.load_tokenizer_from_existing_s3(model_prefix, config)
-
-            logger.info(f"Tokenizer for {model_prefix} not found in S3. Downloading and uploading...")
-            self.download_and_upload_to_s3(model_prefix)
-            logger.info(f"Downloaded and uploaded tokenizer for {model_prefix}. Loading from S3...")
-            return self.load_tokenizer_from_stream(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}")
+
+            tokenizer = self._load_tokenizer(model_prefix, revision)
+            if tokenizer is None:
+                logger.error(f"Failed to load tokenizer for {model_prefix}")
+                raise ValueError(f"Failed to load tokenizer for {model_prefix}")
+            return tokenizer
         except HTTPException as e:
             logger.error(f"Error loading tokenizer: {e}")
             return None
+        except Exception as e:
+            logger.exception(f"Unexpected error loading tokenizer: {e}")
+            raise HTTPException(status_code=500, detail=f"An unexpected error occurred while loading the tokenizer.")
 
-    def load_tokenizer_from_existing_s3(self, model_prefix, config):
-        logger.info(f"Loading tokenizer from S3 for {model_prefix}...")
-        tokenizer_stream = self.stream_from_s3(f"{model_prefix}/tokenizer.json")
-        tokenizer = AutoTokenizer.from_pretrained(None, config=config)
-        logger.info(f"Tokenizer loaded for {model_prefix}.")
-        return tokenizer
-
-    def download_and_upload_to_s3(self, model_prefix, revision="main"):
-        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")
-
-        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")
-
-        logger.info(f"Finished downloading and uploading model files for {model_prefix}.")
+    def _load_tokenizer(self, model_prefix, revision):
+        try:
+            logger.info(f"Downloading tokenizer for {model_prefix} (revision {revision})...")
+            tokenizer_path = hf_hub_download(repo_id=model_prefix, filename="tokenizer.json", revision=revision)
+            return AutoTokenizer.from_pretrained(tokenizer_path)
+        except Exception as e:
+            logger.error(f"Error loading tokenizer: {e}")
+            return None
 
 
     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'))]
+            api = HfApi()
+            model_files = api.list_repo_files(model_prefix, revision=revision)
+            model_files = [file.rfilename for file in model_files if file.rfilename.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}")
+        except Exception as e:
+            logger.error(f"Error retrieving model files 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):