Update app.py

app.py

@@ -1,5 +1,8 @@
 from flask import Flask, request, jsonify
 from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+# Initialize the tokenizer and model
 
 app = Flask(__name__)
 classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
@@ -17,3 +20,27 @@ def classify():
     
     except Exception as e:
         return jsonify({"error": str(e)}), 500
+
+tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
+model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
+nlp = pipeline("ner", model=model, tokenizer=tokenizer)
+@app.route('/ner', methods=['POST'])
+def ner_endpoint():
+    try:
+        # Get text from request
+        data = request.get_json()
+        text = data.get("text", "")
+        
+        # Perform NER
+        ner_results = nlp(text)
+        
+        # Extract words and their corresponding entities
+        words_and_entities = [
+            {"word": result['word'], "entity": result['entity']} 
+            for result in ner_results
+        ]
+        
+        # Return JSON response with the words and their entities
+        return jsonify({"entities": words_and_entities})
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500