Update app.py

app.py

@@ -1,5 +1,6 @@
 import os
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
 import spacy
 import re
 
@@ -7,8 +8,8 @@ import re
 os.environ['TRANSFORMERS_CACHE'] = '/tmp/transformers_cache'
 os.environ['MPLCONFIGDIR'] = '/tmp/.matplotlib'
 
-# Initialize Flask app
-app = Flask(__name__)
+# Initialize FastAPI app
+app = FastAPI()
 
 # Load the spaCy models once
 nlp = spacy.load("en_core_web_sm")
@@ -16,6 +17,10 @@ nlp_coref = spacy.load("en_coreference_web_trf")
 
 REPLACE_PRONOUNS = {"he", "she", "they", "He", "She", "They"}
 
+class CorefRequest(BaseModel):
+    text: str
+    main_characters: str
+
 def extract_core_name(mention_text, main_characters):
     words = mention_text.split()
     for character in main_characters:
@@ -84,14 +89,13 @@ def process_text(text, main_characters):
     else:
         return text
 
-# API endpoint to handle coreference resolution
-@app.route('/predict', methods=['POST'])
-def predict():
-    data = request.json
-    text = data.get('text')
-    main_characters = data.get('main_characters')
-    resolved_text = process_text(text, main_characters.split(","))
-    return jsonify({"resolved_text": resolved_text})
+@app.post("/predict")
+async def predict(coref_request: CorefRequest):
+    resolved_text = process_text(coref_request.text, coref_request.main_characters.split(","))
+    if resolved_text:
+        return {"resolved_text": resolved_text}
+    raise HTTPException(status_code=400, detail="Coreference resolution failed")
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=int(os.getenv("PORT", 7860)))
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))