Update app.py

app.py

@@ -1,4 +1,4 @@
-from huggingface_hub import HfApi, hf_hub_download
+from huggingface_hub import HfApi
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel, field_validator
 import requests
@@ -77,109 +77,109 @@ class S3DirectStream:
     def load_model_from_stream(self, model_prefix):
         try:
             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}")
-
-            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
-
+            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)):
+                logger.info(f"Model {model_prefix} found in S3. Loading...")
+                return self.load_model_from_existing_s3(model_prefix)
+
+            logger.info(f"Model {model_prefix} not found in S3. Downloading and uploading...")
+            self.download_and_upload_to_s3(model_prefix)
+            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}")
-            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_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
 
-    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
+    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}")
 
-        except Exception as e:
-            logger.exception(f"Error loading model: {e}")
-            return None
+                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)
+        return model
 
     def load_tokenizer_from_stream(self, model_prefix):
         try:
             logger.info(f"Loading tokenizer for {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
+            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)
         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(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 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 _get_model_files(self, model_prefix, revision="main"):
+        index_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/"
         try:
-            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'))]
+            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 Exception as e:
-            logger.error(f"Error retrieving model files from Hugging Face: {e}")
+        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}")
             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):