Update app.py

app.py

@@ -1,6 +1,6 @@
 from huggingface_hub import HfApi
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 import requests
 import boto3
 from dotenv import load_dotenv
@@ -35,11 +35,17 @@ s3_client = boto3.client(
 app = FastAPI()
 
 class DownloadModelRequest(BaseModel):
-    model_name: str
+    model_id: str
     pipeline_task: str
     input_text: str
     revision: str = "main"
 
+    @field_validator('model_id')
+    def validate_model_id(cls, value):
+        if not value:
+            raise ValueError("model_id cannot be empty")
+        return value
+
 class S3DirectStream:
     def __init__(self, bucket_name):
         self.s3_client = boto3.client(
@@ -89,7 +95,7 @@ class S3DirectStream:
         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_dict(config_dict)
+        config = AutoConfig.from_pretrained(model_prefix, **config_dict)
         logger.info(f"Config loaded for {model_prefix}.")
 
         if self.file_exists_in_s3(f"{model_prefix}/model.safetensors"):
@@ -115,7 +121,7 @@ class S3DirectStream:
             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) # Pass config
+                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)
@@ -125,16 +131,13 @@ class S3DirectStream:
             logger.error(f"Error loading tokenizer: {e}")
             return None
 
-
-    def load_tokenizer_from_existing_s3(self, model_prefix, config): # Recieve config
+    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) # Use config
+        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"):
         model_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/pytorch_model.bin"
         safetensors_url = f"https://huggingface.co/{model_prefix}/resolve/{revision}/model.safetensors"
@@ -148,7 +151,6 @@ class S3DirectStream:
         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):
         logger.info(f"Downloading from {url}...")
         with requests.get(url, stream=True) as response:
@@ -163,7 +165,6 @@ class S3DirectStream:
             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}")
@@ -172,33 +173,40 @@ class S3DirectStream:
 @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
+        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 # Fixed: Use model_id, not model_name
         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_name, revision)
-        tokenizer = streamer.load_tokenizer_from_stream(model_name)  # Moved after model loading
+        model = streamer.load_model_from_stream(model_id, revision)  # Use model_id
+
+        if model is None:
+            logger.error(f"Failed to load model {model_id}")
+            raise HTTPException(status_code=500, detail=f"Failed to load model {model_id}")
+
+        tokenizer = streamer.load_tokenizer_from_stream(model_id) # Use model_id
+
         logger.info("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("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)
+            inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
             generation_kwargs = dict(inputs, streamer=text_streamer)
             model.generate(**generation_kwargs)
             logger.info("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)
+            outputs = nlp_pipeline(input_text)
             logger.info(f"Pipeline task {task} finished.")
             return {"result": outputs}
 
@@ -206,6 +214,5 @@ async def predict(model_request: DownloadModelRequest):
         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)