Update app.py

app.py

@@ -1,14 +1,18 @@
 import os
 from fastapi import FastAPI, HTTPException, Depends
 from fastapi.responses import JSONResponse
-from pydantic import BaseModel, field_validator
-from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList, pipeline
+from pydantic import BaseModel, field_validator, ValidationError
+from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, GenerationConfig, StoppingCriteriaList, pipeline, StoppingCriteria
 import boto3
 import uvicorn
 import soundfile as sf
 import imageio
-from typing import Dict, Optional
+from typing import Dict, Optional, List
 import torch  # Import torch
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -41,7 +45,7 @@ class GenerateRequest(BaseModel):
     repetition_penalty: float = 1.1
     num_return_sequences: int = 1
     do_sample: bool = True
-    stop_sequences: list[str] = []
+    stop_sequences: List[str] = []
     no_repeat_ngram_size: int = 2
     continuation_id: Optional[str] = None
 
@@ -84,6 +88,7 @@ class S3ModelLoader:
                 tokenizer.pad_token_id = tokenizer.eos_token_id
             return model, tokenizer
         except Exception as e:
+            logging.error(f"Error loading model from S3: {e}")
             raise HTTPException(status_code=500, detail=f"Error loading model from S3: {e}")
 
 model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
@@ -94,6 +99,7 @@ async def get_model_and_tokenizer(model_name: str):
     try:
         return await model_loader.load_model_and_tokenizer(model_name)
     except Exception as e:
+        logging.error(f"Error loading model: {e}")
         raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
 @app.post("/generate")
@@ -142,6 +148,7 @@ async def generate(request: GenerateRequest, model_resources: tuple = Depends(ge
     except HTTPException as http_err:
         raise http_err
     except Exception as e:
+        logging.error(f"Internal server error: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 def generate_text_internal(model, tokenizer, input_text, generation_config, stop_sequences):
@@ -152,6 +159,7 @@ def generate_text_internal(model, tokenizer, input_text, generation_config, stop
 
     class CustomStoppingCriteria(StoppingCriteria): # Inherit directly from StoppingCriteria
         def __init__(self, stop_sequences, tokenizer):
+            super().__init__() # call parent constructor
             self.stop_sequences = stop_sequences
             self.tokenizer = tokenizer
 
@@ -162,7 +170,8 @@ def generate_text_internal(model, tokenizer, input_text, generation_config, stop
                     return True
             return False
 
-    stopping_criteria.append(CustomStoppingCriteria(stop_sequences, tokenizer))
+    if stop_sequences: # Only add if stop_sequences is not empty
+        stopping_criteria.append(CustomStoppingCriteria(stop_sequences, tokenizer))
 
     outputs = model.generate(
         encoded_input.input_ids,
@@ -179,7 +188,8 @@ async def load_pipeline_from_s3(task, model_name):
     try:
         return pipeline(task, model=s3_uri, token=HUGGINGFACE_HUB_TOKEN) # Include token if needed
     except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error loading {task} model from S3: {e}")
+         logging.error(f"Error loading {task} model from S3: {e}")
+         raise HTTPException(status_code=500, detail=f"Error loading {task} model from S3: {e}")
 
 @app.post("/generate-image")
 async def generate_image(request: GenerateRequest):
@@ -198,6 +208,7 @@ async def generate_image(request: GenerateRequest):
     except HTTPException as http_err:
         raise http_err
     except Exception as e:
+        logging.error(f"Internal server error: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 @app.post("/generate-text-to-speech")
@@ -217,6 +228,7 @@ async def generate_text_to_speech(request: GenerateRequest):
     except HTTPException as http_err:
         raise http_err
     except Exception as e:
+        logging.error(f"Internal server error: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
 @app.post("/generate-video")
@@ -236,7 +248,14 @@ async def generate_video(request: GenerateRequest):
     except HTTPException as http_err:
         raise http_err
     except Exception as e:
+        logging.error(f"Internal server error: {str(e)}")
         raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}")
 
+# Adding exception handling for Pydantic validation
+@app.exception_handler(ValidationError)
+async def validation_exception_handler(request, exc):
+    logging.error(f"Validation Error: {exc}")
+    return JSONResponse({"detail": exc.errors()}, status_code=422)
+
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)