Update app.py

app.py

@@ -5,28 +5,29 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 app = FastAPI()
 app.add_middleware(
-  CORSMiddleware, allow_origins=["*"], allow_credentials=False,
-  allow_methods=["*"], allow_headers=["*"]
+    CORSMiddleware,
+    allow_origins=["*"], allow_credentials=False,
+    allow_methods=["*"], allow_headers=["*"]
 )
 
 model_name = "NeuralNovel/Mistral-7B-Instruct-v0.2-Neural-Story"
 tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name, trust_remote_code=True)
 
-class QueryRequest(BaseModel):
-  prompt: str
+class PromptRequest(BaseModel):
+    prompt: str
 
 @app.post("/api/generate-story")
-def generate_story(req: QueryRequest):
-  if not req.prompt.strip():
-    raise HTTPException(status_code=400, detail="Prompt must not be empty")
-  inputs = tokenizer(req.prompt, return_tensors="pt", truncation=True)
-  outputs = model.generate(
-    **inputs,
-    max_new_tokens=200,
-    temperature=0.9,
-    top_p=0.95,
-    repetition_penalty=1.2,
-    do_sample=True
-  )
-  return {"story": tokenizer.decode(outputs[0], skip_special_tokens=True)}
+async def generate_story(req: PromptRequest):
+    if not req.prompt.strip():
+        raise HTTPException(status_code=400, detail="Prompt must not be empty")
+    inputs = tokenizer(req.prompt, return_tensors="pt", truncation=True)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=250,
+        temperature=0.9,
+        top_p=0.95,
+        repetition_penalty=1.2,
+        do_sample=True
+    )
+    return {"story": tokenizer.decode(outputs[0], skip_special_tokens=True)}