Update main.py

main.py

@@ -1,10 +1,8 @@
 from typing import Union
-
 from fastapi import FastAPI
-
 from huggingface_hub import InferenceClient
-
 import os 
+import json
 
 app = FastAPI()
 
@@ -13,11 +11,11 @@ async def root():
     return {"message": "Hello World"}
 
 @app.post("/chat")
-async def chat(prompt: Union[str, None] = None, max_token: Union[str, None] = None):
+async def chat(message: dict):
     try:
         client = InferenceClient(model="gpt2", token=os.environ.get(api_key))
-        input_text = prompt or "Once upon a time"
-        response = client.text_generation(input_text, max_new_tokens=int(max_token or "140"))
+        input_text = message
+        response = client.text_generation(input_text)
         print(response)
         return {"message": response}
     except Exception as e: