Update main.py

main.py

@@ -2,16 +2,20 @@ from fastapi import FastAPI
 from app import predict
 import os
 from huggingface_hub import login
+from pydantic import BaseModel
 
 os.environ['HF_HOME'] = '/hug/cache/'
 os.environ['TRANSFORMERS_CACHE'] = '/blabla/cache/'
 
 app = FastAPI()
 
+class Item(BaseModel):
+    code: str
+
 @app.get("/")
 async def root():
  return {"Code Review Automation":"Version 1.0 'First Draft'"}
 
 @app.post("/AutomateReview/")
-async def predict(input_json: str):
-    return predict(input_json)
+async def predict(item: Item):
+    return predict(item.dict())