|
import streamlit as st |
|
import pandas as pd |
|
import re |
|
import json |
|
import transformers |
|
import torch |
|
from transformers import AutoTokenizer, AutoModelForTokenClassification, Trainer |
|
|
|
st.set_page_config( |
|
page_title="Named Entity Recognition Wolof", |
|
page_icon="📘" |
|
) |
|
|
|
def convert_df(df: pd.DataFrame): |
|
return df.to_csv(index=False).encode('utf-8') |
|
|
|
def convert_json(df: pd.DataFrame): |
|
result = df.to_json(orient="index") |
|
parsed = json.loads(result) |
|
json_string = json.dumps(parsed) |
|
return json_string |
|
|
|
def load_model(): |
|
model = AutoModelForTokenClassification.from_pretrained("vonewman/wolof-finetuned-ner") |
|
trainer = Trainer(model=model) |
|
tokenizer = AutoTokenizer.from_pretrained("vonewman/wolof-finetuned-ner") |
|
return trainer, model, tokenizer |
|
|
|
def align_word_ids(texts): |
|
trainer, model, tokenizer = load_model() |
|
tokenized_inputs = tokenizer(texts, padding='max_length', max_length=218, truncation=True) |
|
word_ids = tokenized_inputs.word_ids() |
|
previous_word_idx = None |
|
label_ids = [] |
|
for word_idx in word_ids: |
|
if word_idx is None: |
|
label_ids.append(-100) |
|
elif word_idx != previous_word_idx: |
|
try: |
|
label_ids.append(1) |
|
except: |
|
label_ids.append(-100) |
|
else: |
|
try: |
|
label_ids.append(1 if label_all_tokens else -100) |
|
except: |
|
label_ids.append(-100) |
|
previous_word_idx = word_idx |
|
return label_ids |
|
|
|
def predict_ner_labels(model, tokenizer, sentence): |
|
use_cuda = torch.cuda.is_available() |
|
device = torch.device("cuda" if use_cuda else "cpu") |
|
if use_cuda: |
|
model = model.cuda() |
|
text = tokenizer(sentence, padding='max_length', max_length=218, truncation=True, return_tensors="pt") |
|
mask = text['attention_mask'].to(device) |
|
input_id = text['input_ids'].to(device) |
|
label_ids = torch.Tensor(align_word_ids(sentence)).unsqueeze(0).to(device) |
|
logits = model(input_id, mask, None) |
|
logits_clean = logits[0][label_ids != -100] |
|
predictions = logits_clean.argmax(dim=1).tolist() |
|
prediction_label = [id2tag[i] for i in predictions] |
|
return prediction_label |
|
|
|
id2tag = {0: 'O', 1: 'B-LOC', 2: 'B-PER', 3: 'I-PER', 4: 'B-ORG', 5: 'I-DATE', 6: 'B-DATE', 7: 'I-ORG', 8: 'I-LOC'} |
|
|
|
def tag_sentence(text): |
|
trainer, model, tokenizer = load_model() |
|
predictions = predict_ner_labels(model, tokenizer, text) |
|
|
|
df = pd.DataFrame({'words': text.split(), 'tags': predictions}) |
|
return df |
|
|
|
st.title("📘 Named Entity Recognition Wolof") |
|
|
|
with st.form(key='my_form'): |
|
x1 = st.text_input(label='Enter a sentence:', max_chars=250) |
|
submit_button = st.form_submit_button(label='🏷️ Create tags') |
|
|
|
if submit_button: |
|
if re.sub('\s+', '', x1) == '': |
|
st.error('Please enter a non-empty sentence.') |
|
elif re.match(r'\A\s*\w+\s*\Z', x1): |
|
st.error("Please enter a sentence with at least one word") |
|
else: |
|
st.markdown("### Tagged Sentence") |
|
st.header("") |
|
results = tag_sentence(x1) |
|
cs, c1, c2, c3, cLast = st.columns([0.75, 1.5, 1.5, 1.5, 0.75]) |
|
with c1: |
|
csvbutton = st.download_button(label="📥 Download .csv", data=convert_df(results), |
|
file_name="results.csv", mime='text/csv', key='csv') |
|
with c2: |
|
textbutton = st.download_button(label="📥 Download .txt", data=convert_df(results), |
|
file_name="results.text", mime='text/plain', key='text') |
|
with c3: |
|
jsonbutton = st.download_button(label="📥 Download .json", data=convert_json(results), |
|
file_name="results.json", mime='application/json', key='json') |
|
st.header("") |
|
c1, c2, c3 = st.columns([1, 3, 1]) |
|
with c2: |
|
st.table(results[['words', 'tags']]) |
|
|
|
st.header("") |
|
st.header("") |
|
st.header("") |
|
with st.expander("ℹ️ - About this app", expanded=True): |
|
st.write( |
|
""" |
|
- The **Named Entity Recognition Wolof** app is a tool that performs named entity recognition in Wolof. |
|
- The available entities are: *corporation*, *location*, *person*, and *date*. |
|
- The app uses the [XLMRoberta model](https://huggingface.co/xlm-roberta-base), fine-tuned on the [masakhaNER](https://huggingface.co/datasets/masakhane/masakhaner2) dataset. |
|
- The model uses the **byte-level BPE tokenizer**. Each sentence is first tokenized. |
|
""" |
|
) |
|
|
