Update main.py

main.py

@@ -60,6 +60,11 @@ def predict_audio(file_path):
 # ScamText model setup
 scamtext_pipe = pipeline("text-classification", model="phishbot/ScamLLM")
 
+# Input model for scam text inference
+class TextInput(BaseModel):
+    input: str
+
+
 # Initialize FastAPI
 app = FastAPI()
 
@@ -75,7 +80,7 @@ async def deepfake_infer(file: UploadFile = File(...)):
     return predictions
 
 @app.post("/scamtext/infer")
-async def scamtext_infer(data: BaseModel):
+async def scamtext_infer(data: TextInput):
     predictions = scamtext_pipe(data.input)
     return predictions