Update main.py

main.py

@@ -47,22 +47,17 @@ async def parse_raw(request: Request):
 
 # Generate text using the Inference API
 def generate_text(input_text: str) -> str:
-    response = openai.Completion.create(
-        engine="text-davinci-003",
-        prompt=input_text,
-        max_tokens=100,
-        n=1,
-        stop=None,
-        temperature=0.5,
-        top_p=1.0,
-        frequency_penalty=0.0,
-        presence_penalty=0.0
-    )
+    data = {"inputs": input_text}
+    response = requests.post(endpoint, json=data, headers=headers)
 
-    print(response)  # Print the entire response object
+    print(response.json())  # Print the entire response object
+
+    try:
+        return response.json()["generated_text"]
+    except KeyError as e:
+        print(f"KeyError: {e}. Check the response object for the correct key.")
+        return ""
 
-    return response.choices[0].text.strip()
-    
 # Route for generating text
 @app.post("/generate_text")
 async def generate_text_route(data: BaseModel = Depends(parse_raw)):