Update app.py

app.py

@@ -2,7 +2,7 @@ import os
 import json
 import logging
 import boto3
-from fastapi import FastAPI, HTTPException
+from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import JSONResponse
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from huggingface_hub import hf_hub_download
@@ -183,7 +183,7 @@ def continue_generation(input_text, model, tokenizer, max_tokens=MAX_TOKENS):
     return generated_text
 
 @app.post("/generate")
-async def generate_text(model_name: str, input_text: str):
+async def generate_text(model_name: str = Query(...), input_text: str = Query(...)):
     try:
         model_loader = S3DirectStream(S3_BUCKET_NAME)
         model = await model_loader.load_model_from_s3(model_name)
@@ -196,7 +196,7 @@ async def generate_text(model_name: str, input_text: str):
         return {"generated_text": generated_text}
 
     except Exception as e:
-        return JSONResponse(status_code=500, content={"detail": str(e)})
+        raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
     import uvicorn