Update app.py

app.py

@@ -1,16 +1,19 @@
 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, field_validator
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     pipeline,
-    GenerationConfig
+    GenerationConfig,
+    StoppingCriteriaList
 )
 import boto3
 from huggingface_hub import hf_hub_download
@@ -18,9 +21,8 @@ import soundfile as sf
 import numpy as np
 import torch
 import uvicorn
-from tqdm import tqdm
 
-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
@@ -40,6 +42,23 @@ class GenerateRequest(BaseModel):
     repetition_penalty: float = 1.0
     num_return_sequences: int = 1
     do_sample: bool = True
+    chunk_delay: float = 0.0
+    stop_sequences: list[str] = []
+
+    model_config = {"protected_namespaces": ()}
+
+    @field_validator("model_name")
+    def model_name_cannot_be_empty(cls, v):
+        if not v:
+            raise ValueError("model_name cannot be empty.")
+        return v
+
+    @field_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):
@@ -55,15 +74,23 @@ class S3ModelLoader:
             logging.info(f"Trying to load {model_name} from S3...")
             config = AutoConfig.from_pretrained(s3_uri)
             model = AutoModelForCausalLM.from_pretrained(s3_uri, config=config)
-            tokenizer = AutoTokenizer.from_pretrained(s3_uri)
+            tokenizer = AutoTokenizer.from_pretrained(s3_uri, config=config)
+
+            if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
+                tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
+
             logging.info(f"Loaded {model_name} from S3 successfully.")
             return model, tokenizer
         except EnvironmentError:
             logging.info(f"Model {model_name} not found in S3. Downloading...")
             try:
-                with tqdm(unit="B", unit_scale=True, desc=f"Downloading {model_name}") as t:
-                    model = AutoModelForCausalLM.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN, _tqdm=t)
-                    tokenizer = AutoTokenizer.from_pretrained(model_name, token=HUGGINGFACE_HUB_TOKEN)
+                config = AutoConfig.from_pretrained(model_name)
+                tokenizer = AutoTokenizer.from_pretrained(model_name, config=config)
+                model = AutoModelForCausalLM.from_pretrained(model_name, config=config, token=HUGGINGFACE_HUB_TOKEN)
+
+                if tokenizer.eos_token_id is not None and tokenizer.pad_token_id is None:
+                    tokenizer.pad_token_id = config.pad_token_id or tokenizer.eos_token_id
+
                 logging.info(f"Downloaded {model_name} successfully.")
                 logging.info(f"Saving {model_name} to S3...")
                 model.save_pretrained(s3_uri)
@@ -71,7 +98,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,44 +109,49 @@ 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
             )
 
             async def stream_text():
-                input_text = body.input_text
-                max_length = model.config.max_position_embeddings
+                input_text = validated_body.input_text
                 generated_text = ""
+                max_length = model.config.max_position_embeddings
 
                 while True:
-                    inputs = tokenizer(input_text, return_tensors="pt").to(device)
-                    input_length = inputs.input_ids.shape[1]
+                    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
+                    input_length = encoded_input["input_ids"].shape[1]
                     remaining_tokens = max_length - input_length
-                    if remaining_tokens < body.max_new_tokens:
-                        generation_config.max_new_tokens = remaining_tokens
-                        if remaining_tokens <= 0:
-                            break
 
-                    output = model.generate(**inputs, generation_config=generation_config)
+                    if remaining_tokens <= 0:
+                        break
+
+                    generation_config.max_new_tokens = min(remaining_tokens, validated_body.max_new_tokens)
+
+                    stopping_criteria = StoppingCriteriaList(
+                        [lambda _, outputs: tokenizer.decode(outputs[0][-1], skip_special_tokens=True) in validated_body.stop_sequences] if validated_body.stop_sequences else []
+                    )
+
+                    output = model.generate(**encoded_input, generation_config=generation_config, stopping_criteria=stopping_criteria)
                     chunk = tokenizer.decode(output[0], skip_special_tokens=True)
                     generated_text += chunk
                     yield chunk
-                    if len(tokenizer.encode(generated_text)) >= max_length:
-                        break
-                    input_text = chunk
+                    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,24 +159,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 = 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}")
@@ -154,8 +186,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__":