File size: 4,592 Bytes
08eb663 345d423 08eb663 345d423 08eb663 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 |
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)
# Créez un DataFrame avec les colonnes "words" et "tags"
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.
"""
)
|