Update app.py

app.py

@@ -1,25 +1,12 @@
 from pydantic import BaseModel
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, HTTPException
+
 from happytransformer import HappyTextToText, TTSettings
-from fastapi import FastAPI
 
 app = FastAPI()
-class Generate(BaseModel):
-    text:str
-
-def generate_text(prompt: str):
-    if prompt == "":
-        return {"detail": "Please provide a prompt."}
-    else:
-
-
-        happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction")
-        
-        args = TTSettings(num_beams=5, min_length=1,  max_length=100000)
-        
-        result = happy_tt.generate_text(f"grammar: {prompt}", args=args)
-        return result.text
 
+# Enable CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -28,10 +15,23 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
+class GenerateInput(BaseModel):
+    prompt: str
+
+happy_tt = HappyTextToText("T5", "vennify/t5-base-grammar-correction")
+args = TTSettings(num_beams=5, min_length=1, max_length=100000)
+
+def generate_text(prompt: str):
+    if not prompt.strip():
+        raise HTTPException(status_code=400, detail="Please provide a non-empty prompt.")
+    
+    result = happy_tt.generate_text(f"grammar: {prompt}", args=args)
+    return result.text
+
 @app.get("/", tags=["Home"])
 def api_home():
     return {'detail': 'Welcome to FastAPI TextGen Tutorial!'}
 
-@app.post("/api/generate",  tags=["Generate"], response_model=Generate)
-def inference(input_prompt: str):
-    return generate_text(prompt=input_prompt)
+@app.post("/api/generate", tags=["Generate"])
+def inference(data: GenerateInput):
+    return generate_text(data.prompt)