Update app.py

app.py

@@ -2,16 +2,18 @@ import os
 import logging
 import time
 from io import BytesIO
+from typing import Union
 
-from fastapi import FastAPI, HTTPException, Response, Request
+from fastapi import FastAPI, HTTPException, Response, Request, UploadFile, File
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationError, validator
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     pipeline,
-    GenerationConfig
+    GenerationConfig,
+    StoppingCriteriaList
 )
 import boto3
 from huggingface_hub import hf_hub_download
@@ -20,7 +22,7 @@ import numpy as np
 import torch
 import uvicorn
 
-logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(filename)s:%(lineno)d - %(message)s")
 
 AWS_ACCESS_KEY_ID = os.getenv("AWS_ACCESS_KEY_ID")
 AWS_SECRET_ACCESS_KEY = os.getenv("AWS_SECRET_ACCESS_KEY")
@@ -30,7 +32,7 @@ HUGGINGFACE_HUB_TOKEN = os.getenv("HUGGINGFACE_HUB_TOKEN")
 
 class GenerateRequest(BaseModel):
     model_name: str
-    input_text: str
+    input_text: str = ""
     task_type: str
     temperature: float = 1.0
     max_new_tokens: int = 200
@@ -41,6 +43,20 @@ class GenerateRequest(BaseModel):
     num_return_sequences: int = 1
     do_sample: bool = True
     chunk_delay: float = 0.0
+    stop_sequences: list[str] = []
+
+    @validator("model_name")
+    def model_name_cannot_be_empty(cls, v):
+        if not v:
+            raise ValueError("model_name cannot be empty.")
+        return v
+
+    @validator("task_type")
+    def task_type_must_be_valid(cls, v):
+        valid_types = ["text-to-text", "text-to-image", "text-to-speech", "text-to-video"]
+        if v not in valid_types:
+            raise ValueError(f"task_type must be one of: {valid_types}")
+        return v
 
 class S3ModelLoader:
     def __init__(self, bucket_name, s3_client):
@@ -57,6 +73,13 @@ class S3ModelLoader:
             config = AutoConfig.from_pretrained(s3_uri)
             model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config)
             tokenizer = AutoTokenizer.from_pretrained(s3_uri)
+
+            if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
+                if config.pad_token_id is not None:
+                    tokenizer.pad_token_id = config.pad_token_id
+                else:
+                    tokenizer.pad_token_id = 0
+
             logging.info(f"Loaded {model_name} from S3 successfully.")
             return model, tokenizer
         except EnvironmentError:
@@ -64,6 +87,14 @@ class S3ModelLoader:
             try:
                 model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
                 tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
+                
+                if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
+                    config = AutoConfig.from_pretrained(model_name)
+                    if config.pad_token_id is not None:
+                        tokenizer.pad_token_id = config.pad_token_id
+                    else:
+                        tokenizer.pad_token_id = 0
+
                 logging.info(f"Downloaded {model_name} successfully.")
                 logging.info(f"Saving {model_name} to S3...")
                 model.save_pretrained(s3_uri)
@@ -71,7 +102,7 @@ class S3ModelLoader:
                 logging.info(f"Saved {model_name} to S3 successfully.")
                 return model, tokenizer
             except Exception as e:
-                logging.error(f"Error downloading/uploading model: {e}")
+                logging.exception(f"Error downloading/uploading model: {e}")
                 raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
 app = FastAPI()
@@ -82,23 +113,27 @@ model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 @app.post("/generate")
 async def generate(request: Request, body: GenerateRequest):
     try:
-        model, tokenizer = await model_loader.load_model_and_tokenizer(body.model_name)
+        validated_body = GenerateRequest(**body.model_dump())
+        model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
 
-        if body.task_type == "text-to-text":
+        if validated_body.task_type == "text-to-text":
             generation_config = GenerationConfig(
-                temperature=body.temperature,
-                max_new_tokens=body.max_new_tokens,
-                top_p=body.top_p,
-                top_k=body.top_k,
-                repetition_penalty=body.repetition_penalty,
-                do_sample=body.do_sample,
-                num_return_sequences=body.num_return_sequences
+                temperature=validated_body.temperature,
+                max_new_tokens=validated_body.max_new_tokens,
+                top_p=validated_body.top_p,
+                top_k=validated_body.top_k,
+                repetition_penalty=validated_body.repetition_penalty,
+                do_sample=validated_body.do_sample,
+                num_return_sequences=validated_body.num_return_sequences,
+                stopping_criteria=StoppingCriteriaList(
+                    [lambda _, outputs: tokenizer.decode(outputs[0][-1]) in validated_body.stop_sequences] if validated_body.stop_sequences else []
+                )
             )
 
             async def stream_text():
-                input_text = body.input_text
+                input_text = validated_body.input_text
                 generated_text = ""
                 max_length = model.config.max_position_embeddings
 
@@ -110,16 +145,16 @@ async def generate(request: Request, body: GenerateRequest):
                     if remaining_tokens <= 0:
                         break
 
-                    generation_config.max_new_tokens = min(remaining_tokens, body.max_new_tokens)
+                    generation_config.max_new_tokens = min(remaining_tokens, validated_body.max_new_tokens)
 
                     output = model.generate(**encoded_input, generation_config=generation_config)
                     chunk = tokenizer.decode(output[0], skip_special_tokens=True)
                     generated_text += chunk
                     yield chunk
-                    time.sleep(body.chunk_delay)
+                    time.sleep(validated_body.chunk_delay)
                     input_text = generated_text
 
-            if body.stream:
+            if validated_body.stream:
                 return StreamingResponse(stream_text(), media_type="text/plain")
             else:
                 generated_text = ""
@@ -127,32 +162,24 @@ async def generate(request: Request, body: GenerateRequest):
                     generated_text += chunk
                 return {"result": generated_text}
 
-        elif body.task_type == "text-to-image":
+        elif validated_body.task_type == "text-to-image":
             generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=device)
-            image = generator(body.input_text)[0]
+            image = generator(validated_body.input_text)[0]
             image_bytes = image.tobytes()
             return Response(content=image_bytes, media_type="image/png")
 
-        elif body.task_type == "text-to-speech":
+        elif validated_body.task_type == "text-to-speech":
             generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
-            audio = generator(body.input_text)
-            audio_bytesio = BytesIO()
-            sf.write(audio_bytesio, audio["sampling_rate"], np.int16(audio["audio"]))
-            audio_bytes = audio_bytesio.getvalue()
-            return Response(content=audio_bytes, media_type="audio/wav")
-
-        elif body.task_type == "text-to-audio":
-            generator = pipeline("text-to-audio", model=model, tokenizer=tokenizer, device=device)
-            audio = generator(body.input_text)
+            audio = generator(validated_body.input_text)
             audio_bytesio = BytesIO()
             sf.write(audio_bytesio, audio["sampling_rate"], np.int16(audio["audio"]))
             audio_bytes = audio_bytesio.getvalue()
             return Response(content=audio_bytes, media_type="audio/wav")
 
-        elif body.task_type == "text-to-video":
+        elif validated_body.task_type == "text-to-video":
             try:
                 generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
-                video = generator(body.input_text)
+                video = generator(validated_body.input_text)
                 return Response(content=video, media_type="video/mp4")
             except Exception as e:
                 raise HTTPException(status_code=500, detail=f"Error in text-to-video generation: {e}")
@@ -162,8 +189,11 @@ async def generate(request: Request, body: GenerateRequest):
 
     except HTTPException as e:
         raise e
+    except ValidationError as e:
+        raise HTTPException(status_code=422, detail=e.errors())
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        logging.exception(f"An unexpected error occurred: {e}")
+        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
 
 
 if __name__ == "__main__":