Nechba commited on
Commit
91a018b
·
verified ·
1 Parent(s): b76b429

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +27 -0
app.py CHANGED
@@ -1,5 +1,8 @@
1
  from flask import Flask, request, jsonify
2
  from transformers import pipeline
 
 
 
3
 
4
  app = Flask(__name__)
5
  classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
@@ -17,3 +20,27 @@ def classify():
17
 
18
  except Exception as e:
19
  return jsonify({"error": str(e)}), 500
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
  from flask import Flask, request, jsonify
2
  from transformers import pipeline
3
+ from transformers import AutoTokenizer, AutoModelForTokenClassification
4
+
5
+ # Initialize the tokenizer and model
6
 
7
  app = Flask(__name__)
8
  classifier = pipeline("text-classification", model="j-hartmann/emotion-english-distilroberta-base", return_all_scores=True)
 
20
 
21
  except Exception as e:
22
  return jsonify({"error": str(e)}), 500
23
+
24
+ tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
25
+ model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
26
+ nlp = pipeline("ner", model=model, tokenizer=tokenizer)
27
+ @app.route('/ner', methods=['POST'])
28
+ def ner_endpoint():
29
+ try:
30
+ # Get text from request
31
+ data = request.get_json()
32
+ text = data.get("text", "")
33
+
34
+ # Perform NER
35
+ ner_results = nlp(text)
36
+
37
+ # Extract words and their corresponding entities
38
+ words_and_entities = [
39
+ {"word": result['word'], "entity": result['entity']}
40
+ for result in ner_results
41
+ ]
42
+
43
+ # Return JSON response with the words and their entities
44
+ return jsonify({"entities": words_and_entities})
45
+ except Exception as e:
46
+ return jsonify({"error": str(e)}), 500