cindyangelira commited on
Commit
6d35993
·
1 Parent(s): 66bfec7

Add application file

Browse files
Files changed (3) hide show
  1. .gitignore +2 -0
  2. app.py +60 -0
  3. ner_colored_output.html +0 -0
.gitignore ADDED
@@ -0,0 +1,2 @@
 
 
 
1
+ .env
2
+ demo.ipynb
app.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ from transformers import pipeline
3
+ from spacy import displacy
4
+
5
+ # load model pipeline globally
6
+ ner_pipe = pipeline("token-classification", model="cindyangelira/ner-roberta-large-bahasa-indonesia")
7
+
8
+ # define colors for each tag
9
+ def get_colors():
10
+ return {
11
+ "O": "#ffffff", # White for 'O'
12
+ "PERSON": "#ffadad", # Light red for 'PERSON'
13
+ "LOCATION": "#ffda83", # Light yellow for 'LOCATION'
14
+ "DOB": "#ffadad", # Light red for 'DOB'
15
+ "EMAIL": "#85e0e0", # Light cyan for 'EMAIL'
16
+ "GENDER": "#c3c3e0", # Light gray for 'GENDER'
17
+ "ACCOUNT": "#b0e0e6", # Light blue for 'ACCOUNT'
18
+ "ID": "#800080", # Purple for 'ID'
19
+ "PHONE": "#d1ff85" # Light green for 'PHONE NUMBER'
20
+ }
21
+
22
+
23
+ def process_prediction(text, pred):
24
+ colors = get_colors()
25
+
26
+ for token in pred:
27
+ token['label'] = token['entity'].replace('B-', '').replace('I-', '')
28
+
29
+ ents = [{'start': token['start'], 'end': token['end'], 'label': token['label']} for token in pred]
30
+
31
+ doc = {
32
+ "text": text,
33
+ "ents": ents,
34
+ "title": None
35
+ }
36
+
37
+ options = {"ents": list(colors.keys()), "colors": colors}
38
+ html = displacy.render(doc, style="ent", manual=True, options=options)
39
+ return html
40
+
41
+
42
+ def ner_visualization(text):
43
+ predictions = ner_pipe(text)
44
+ return process_prediction(text, predictions)
45
+
46
+
47
+ def build_interface():
48
+ iface = gr.Interface(
49
+ fn=ner_visualization, # Main function for NER visualization
50
+ inputs=gr.Textbox(label="Input Text"),# Input textbox
51
+ outputs="html", # Output is HTML with rendered NER
52
+ title="NER Bahasa Indonesia", # Title of the app
53
+ description="Enter text to see named entity recognition results highlighted."
54
+ )
55
+ return iface
56
+
57
+
58
+ if __name__ == "__main__":
59
+ app = build_interface()
60
+ app.launch()
ner_colored_output.html ADDED
File without changes