File size: 3,947 Bytes
4d9dd77 cb0039e f5b6e30 cb0039e 4d9dd77 cb0039e f5b6e30 cb0039e 4d9dd77 f5b6e30 4d9dd77 f5b6e30 cb0039e a055f19 cb0039e a055f19 cb0039e fe3337c 52d47ae 18ca43a fe3337c 18ca43a fe3337c 18ca43a fe3337c f5b6e30 cb0039e f5b6e30 3f67191 7c44018 e12504f 67c7e19 f5b6e30 cb0039e f5b6e30 cb0039e f5b6e30 cb0039e f5b6e30 cb0039e f5b6e30 cb0039e f5b6e30 7c44018 |
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 |
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})
df['tags'] = df['tags'].map(lambda x: f'background-color: lightblue' if x != 'O' else '')
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
|