Update app.py

app.py

@@ -10,9 +10,14 @@ from transformers import pipeline, AutoModelForCausalLM, AutoTokenizer, AutoConf
 import safetensors.torch
 import torch
 from fastapi.responses import StreamingResponse
+from tqdm import tqdm
+import logging
 
 load_dotenv()
 
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+logger = logging.getLogger(__name__)
+
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
 AWS_REGION = os.getenv("AWS_REGION")
@@ -46,60 +51,83 @@ class S3DirectStream:
 
     def stream_from_s3(self, key):
         try:
+            logger.info(f"Downloading {key} from S3...")
             response = self.s3_client.get_object(Bucket=self.bucket_name, Key=key)
+            logger.info(f"Downloaded {key} from S3 successfully.")
             return response['Body']
         except self.s3_client.exceptions.NoSuchKey:
+            logger.error(f"File {key} not found in S3")
             raise HTTPException(status_code=404, detail=f"File {key} not found in S3")
 
     def file_exists_in_s3(self, key):
         try:
             self.s3_client.head_object(Bucket=self.bucket_name, Key=key)
+            logger.info(f"File {key} exists in S3.")
             return True
         except self.s3_client.exceptions.ClientError:
+            logger.info(f"File {key} does not exist in S3.")
             return False
 
     def load_model_from_stream(self, model_prefix, revision):
         try:
+            logger.info(f"Loading model {model_prefix} (revision {revision})...")
             if self.file_exists_in_s3(f"{model_prefix}/config.json") and \
                (self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin") or self.file_exists_in_s3(f"{model_prefix}/model.safetensors")):
+                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, revision)
+            logger.info(f"Downloaded and uploaded {model_prefix}. Loading from S3...")
             return self.load_model_from_stream(model_prefix, revision)
         except HTTPException as e:
+            logger.error(f"Error loading model: {e}")
             return None
 
     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 = AutoConfig.from_pretrained(config_stream) # Directly from stream
+        config = AutoConfig.from_pretrained(config_stream)
+        logger.info(f"Config loaded for {model_prefix}.")
 
         if self.file_exists_in_s3(f"{model_prefix}/model.safetensors"):
+            logger.info(f"Loading safetensors model for {model_prefix} from S3...")
             model_stream = self.stream_from_s3(f"{model_prefix}/model.safetensors")
             model = AutoModelForCausalLM.from_config(config)
             model.load_state_dict(safetensors.torch.load_stream(model_stream))
+            logger.info(f"Safetensors model loaded for {model_prefix}.")
         elif self.file_exists_in_s3(f"{model_prefix}/pytorch_model.bin"):
+            logger.info(f"Loading PyTorch model for {model_prefix} from S3...")
             model_stream = self.stream_from_s3(f"{model_prefix}/pytorch_model.bin")
             model = AutoModelForCausalLM.from_config(config)
-            state_dict = torch.load(model_stream, map_location="cpu")  # Load directly
+            state_dict = torch.load(model_stream, map_location="cpu")
             model.load_state_dict(state_dict)
+            logger.info(f"PyTorch model loaded for {model_prefix}.")
         else:
-           raise EnvironmentError(f"No model file found for {model_prefix} in S3")
+            logger.error(f"No model file found for {model_prefix} in S3")
+            raise EnvironmentError(f"No model file found for {model_prefix} in S3")
         return model
 
-
-
     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)
+
+            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
 
     def load_tokenizer_from_existing_s3(self, model_prefix):
+        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(tokenizer_stream) # Directly from stream
+        tokenizer = AutoTokenizer.from_pretrained(tokenizer_stream)
+        logger.info(f"Tokenizer loaded for {model_prefix}.")
         return tokenizer
 
 
@@ -109,56 +137,69 @@ class S3DirectStream:
         tokenizer_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/tokenizer.json"
         config_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/config.json"
 
+        logger.info(f"Downloading and uploading model files for {model_prefix} to S3...")
         self.download_and_upload_to_s3_url(model_url, f"{model_prefix}/pytorch_model.bin")
         self.download_and_upload_to_s3_url(safetensors_url, f"{model_prefix}/model.safetensors")
         self.download_and_upload_to_s3_url(tokenizer_url, f"{model_prefix}/tokenizer.json")
         self.download_and_upload_to_s3_url(config_url, f"{model_prefix}/config.json")
-
+        logger.info(f"Finished downloading and uploading model files for {model_prefix}.")
 
     def download_and_upload_to_s3_url(self, url, s3_key):
-        response = requests.get(url, stream=True)
-        if response.status_code == 200:
-            self.s3_client.upload_fileobj(response.raw, self.bucket_name, s3_key) # Direct upload
-        elif response.status_code == 404:
-            raise HTTPException(status_code=404, detail=f"Error downloading file from {url}. File not found.")
-
-        else:
-            raise HTTPException(status_code=500, detail=f"Error downloading file from {url}")
-
+        logger.info(f"Downloading from {url}...")
+        with requests.get(url, stream=True) as response:
+            if response.status_code == 200:
+                total_size_in_bytes= int(response.headers.get('content-length', 0))
+                block_size = 1024
+                progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+                logger.info(f"Uploading to S3: {s3_key}...")
+                self.s3_client.upload_fileobj(response.raw, self.bucket_name, s3_key, Callback=lambda bytes_transferred: progress_bar.update(bytes_transferred))
+                progress_bar.close()
+                logger.info(f"Uploaded {s3_key} to S3 successfully.")
+            elif response.status_code == 404:
+                logger.error(f"File not found at {url}")
+                raise HTTPException(status_code=404, detail=f"Error downloading file from {url}. File not found.")
+            else:
+                logger.error(f"Error downloading from {url}: Status code {response.status_code}")
+                raise HTTPException(status_code=500, detail=f"Error downloading file from {url}")
 
 
 @app.post("/predict/")
 async def predict(model_request: DownloadModelRequest):
     try:
+        logger.info(f"Received request: Model={model_request.model_name}, Task={model_request.pipeline_task}, Input={model_request.input_text}")
         model_name = model_request.model_name
         revision = model_request.revision
 
         streamer = S3DirectStream(S3_BUCKET_NAME)
+        logger.info(f"Loading model and tokenizer...")
         model = streamer.load_model_from_stream(model_name, revision)
         tokenizer = streamer.load_tokenizer_from_stream(model_name)
+        logger.info(f"Model and tokenizer loaded.")
 
         task = model_request.pipeline_task
         if task not in ["text-generation", "sentiment-analysis", "translation", "fill-mask", "question-answering",  "summarization", "zero-shot-classification"]:
             raise HTTPException(status_code=400, detail="Unsupported pipeline task")
 
         if task == "text-generation":
+            logger.info(f"Starting text generation...")
             text_streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
             inputs = tokenizer(model_request.input_text, return_tensors="pt").to(model.device)
             generation_kwargs = dict(inputs, streamer=text_streamer)
             model.generate(**generation_kwargs)
+            logger.info(f"Text generation finished.")
             return StreamingResponse(iter([tokenizer.decode(token) for token in text_streamer]), media_type="text/event-stream")
 
         else:
+            logger.info(f"Starting pipeline task: {task}...")
             nlp_pipeline = pipeline(task, model=model, tokenizer=tokenizer, device_map="auto", trust_remote_code=True)
             outputs = nlp_pipeline(model_request.input_text)
+            logger.info(f"Pipeline task {task} finished.")
             return {"result": outputs}
 
-
     except Exception as e:
-        print(f"Complete Error: {e}")
+        logger.exception(f"Error processing request: {e}")
         raise HTTPException(status_code=500, detail=f"Error processing request: {str(e)}")
 
 
-
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)