add token limit

main.py

@@ -7,7 +7,9 @@ import torch
 
 app = FastAPI()
 
-async def generate_text(title, max_length=1000, top_k=50, model_dir="./model/custom-gpt2-model", tokenizer_dir="./model/custom-gpt2-tokenizer"):
+async def generate_text(title, max_length=400, top_k=50, model_dir="./model/custom-gpt2-model", tokenizer_dir="./model/custom-gpt2-tokenizer"):
+    if max_length > 400:
+        return {"error": "Limit for max_length is 400 tokens."}
     model = AutoModelForCausalLM.from_pretrained(model_dir)
     tokenizer = AutoTokenizer.from_pretrained(tokenizer_dir)
     input_text = f"[TITLE] {title} [/TITLE]"
@@ -25,7 +27,7 @@ async def generate_text(title, max_length=1000, top_k=50, model_dir="./model/cus
 
 class RequestParams(BaseModel):
     title: str
-    max_length: int = 1000
+    max_length: int = 400
     top_k: int = 50
 
 @app.post("/generate-article")