Update app.py

app.py

@@ -1,19 +1,16 @@
 import os
 import logging
-import time
 from io import BytesIO
-from typing import Union
 
-from fastapi import FastAPI, HTTPException, Response, Request, UploadFile, File
+from fastapi import FastAPI, HTTPException, Response, Request
 from fastapi.responses import StreamingResponse
-from pydantic import BaseModel, ValidationError, field_validator
+from pydantic import BaseModel
 from transformers import (
     AutoConfig,
     AutoModelForCausalLM,
     AutoTokenizer,
     pipeline,
-    GenerationConfig,
-    StoppingCriteriaList
+    GenerationConfig
 )
 import boto3
 from huggingface_hub import hf_hub_download
@@ -21,8 +18,9 @@ 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 - %(filename)s:%(lineno)d - %(message)s")
+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")
@@ -32,7 +30,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
@@ -42,23 +40,6 @@ 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):
@@ -74,23 +55,15 @@ 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, 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
-
+            tokenizer = AutoTokenizer.from_pretrained(s3_uri)
             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:
-                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
-
+                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)
                 logging.info(f"Downloaded {model_name} successfully.")
                 logging.info(f"Saving {model_name} to S3...")
                 model.save_pretrained(s3_uri)
@@ -98,7 +71,7 @@ class S3ModelLoader:
                 logging.info(f"Saved {model_name} to S3 successfully.")
                 return model, tokenizer
             except Exception as e:
-                logging.exception(f"Error downloading/uploading model: {e}")
+                logging.error(f"Error downloading/uploading model: {e}")
                 raise HTTPException(status_code=500, detail=f"Error loading model: {e}")
 
 app = FastAPI()
@@ -109,49 +82,44 @@ model_loader = S3ModelLoader(S3_BUCKET_NAME, s3_client)
 @app.post("/generate")
 async def generate(request: Request, body: GenerateRequest):
     try:
-        validated_body = GenerateRequest(**body.model_dump())
-        model, tokenizer = await model_loader.load_model_and_tokenizer(validated_body.model_name)
+        model, tokenizer = await model_loader.load_model_and_tokenizer(body.model_name)
         device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
 
-        if validated_body.task_type == "text-to-text":
+        if body.task_type == "text-to-text":
             generation_config = GenerationConfig(
-                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
+                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
             )
 
             async def stream_text():
-                input_text = validated_body.input_text
-                generated_text = ""
+                input_text = body.input_text
                 max_length = model.config.max_position_embeddings
+                generated_text = ""
 
                 while True:
-                    encoded_input = tokenizer(input_text, return_tensors="pt", truncation=True, max_length=max_length).to(device)
-                    input_length = encoded_input["input_ids"].shape[1]
+                    inputs = tokenizer(input_text, return_tensors="pt").to(device)
+                    input_length = inputs.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
 
-                    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)
+                    output = model.generate(**inputs, generation_config=generation_config)
                     chunk = tokenizer.decode(output[0], skip_special_tokens=True)
                     generated_text += chunk
                     yield chunk
-                    time.sleep(validated_body.chunk_delay)
-                    input_text = generated_text
+                    if len(tokenizer.encode(generated_text)) >= max_length:
+                        break
+                    input_text = chunk
 
-            if validated_body.stream:
+            if body.stream:
                 return StreamingResponse(stream_text(), media_type="text/plain")
             else:
                 generated_text = ""
@@ -159,24 +127,24 @@ async def generate(request: Request, body: GenerateRequest):
                     generated_text += chunk
                 return {"result": generated_text}
 
-        elif validated_body.task_type == "text-to-image":
+        elif body.task_type == "text-to-image":
             generator = pipeline("text-to-image", model=model, tokenizer=tokenizer, device=device)
-            image = generator(validated_body.input_text)[0]
+            image = generator(body.input_text)[0]
             image_bytes = image.tobytes()
             return Response(content=image_bytes, media_type="image/png")
 
-        elif validated_body.task_type == "text-to-speech":
+        elif body.task_type == "text-to-speech":
             generator = pipeline("text-to-speech", model=model, tokenizer=tokenizer, device=device)
-            audio = generator(validated_body.input_text)
+            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 validated_body.task_type == "text-to-video":
+        elif body.task_type == "text-to-video":
             try:
                 generator = pipeline("text-to-video", model=model, tokenizer=tokenizer, device=device)
-                video = generator(validated_body.input_text)
+                video = generator(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}")
@@ -186,12 +154,9 @@ 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:
-        logging.exception(f"An unexpected error occurred: {e}")
-        raise HTTPException(status_code=500, detail="An unexpected error occurred.")
+        raise HTTPException(status_code=500, detail=str(e))
 
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=8000)