Updates

app.py

@@ -51,14 +51,36 @@ print("Start loading")
 download_folder_from_hf("SinaLab/ArabGlossBERT", "bert-base-arabertv02_22_May_2021_00h_allglosses_unused01")
 download_folder_from_hf("SinaLab/ArabGlossBERT", "bert-base-arabertv02")   
 download_file_from_hf("SinaLab/ArabGlossBERT","one_gram.pickle")
+download_file_from_hf("SinaLab/ArabGlossBERT","two_grams.pickle")
+download_file_from_hf("SinaLab/ArabGlossBERT","three_grams.pickle")
+download_file_from_hf("SinaLab/ArabGlossBERT","four_grams.pickle")
+download_file_from_hf("SinaLab/ArabGlossBERT","five_grams.pickle")
+download_file_from_hf("SinaLab/ALMA","lemmas_dic.pickle")
 
 print("Finish loading")
 
 from sinatools.wsd.disambiguator import disambiguate
+from pydantic import BaseModel
+from fastapi.responses import JSONResponse
 
-@app.get("/predict")
-def greet_json():
+class SALMARequest(BaseModel):
+    text: str
+
+@app.post("/predict")
+def predict(request: SALMARequest):
+    # Load tagger
+    text = request.text  
+    
     print("Start disambiguate")
-    salma_output = disambiguate('تمشيت بين الجداول والأنهار')
-    print("End disambiguate")
-    return {"output": salma_output}
+    salma_output = disambiguate(text)
+    content = {
+        "resp": salma_output,
+        "statusText": "OK",
+        "statusCode": 0,
+    }
+
+    return JSONResponse(
+        content=content,
+        media_type="application/json",
+        status_code=200,
+    )