Update app.py

app.py

@@ -26,10 +26,38 @@ nlp = pipeline("ner", model=model_name, tokenizer=model_name)
 class TextRequest(BaseModel):
     text: str
 
-@app.post("/predict/")
+# Define a Pydantic model for the response
+class Entity(BaseModel):
+    entity: str
+    score: float
+    index: int
+    word: str
+    start: int
+    end: int
+
+
+@app.post("/predict/", response_model=List[Entity])
 def predict(request: TextRequest):
-    ner_results = nlp(request.text)
-    return {"entities":ner_results}
+    try:
+        # Use the model to perform NER on the input text
+        ner_results = nlp(request.text)
+        
+        # Format the results to match the desired output structure
+        entities = [
+            {
+                "entity": result["entity"],
+                "score": result["score"],
+                "index": result["index"],
+                "word": result["word"],
+                "start": result["start"],
+                "end": result["end"]
+            }
+            for result in ner_results
+        ]
+        
+        return entities
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
 
 
 @app.get("/")