Kabeer Akande commited on
Commit
3d1d2fa
1 Parent(s): 532c912

adds model inference script

Browse files
Files changed (1) hide show
  1. app.py +31 -0
app.py ADDED
@@ -0,0 +1,31 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import pipeline
2
+ import gradio as gr
3
+
4
+ #reference appropriate Hugging Face model
5
+ model_name = 'koakande/bert-finetuned-ner'
6
+
7
+ # Load token classification pipeline modelfrom Hugging Face
8
+ model = pipeline("token-classification", model=model_name, aggregation_strategy="simple")
9
+
10
+ # write a prediction method for the model
11
+ def predict_entities(text):
12
+ # Use the loaded model to identify entities in the text
13
+ entities = model(text)
14
+ # Highlight identified entities in the input text
15
+ highlighted_text = text
16
+ for entity in entities:
17
+ entity_text = text[entity['start']:entity['end']]
18
+ replacement = f"<span style='border: 2px solid green;'>{entity_text}</span>"
19
+ highlighted_text = highlighted_text.replace(entity_text, replacement)
20
+ return highlighted_text
21
+
22
+ # gradio interface
23
+ iface = gr.Interface(
24
+ fn=predict_entities,
25
+ inputs=gr.Textbox(lines=5, placeholder="Enter text..."),
26
+ outputs=gr.HTML(),
27
+ title="Named Entity Identification",
28
+ description="Enter text to identify entities using the model.",
29
+ )
30
+
31
+ iface.launch()