jcg00v's picture
Update app.py
73b559a verified
raw
history blame
5.71 kB
import streamlit as st
from PIL import Image
from transformers import pipeline, AutoModelForTokenClassification, AutoTokenizer
def clear_text():
st.session_state.my_text = st.session_state.widget
st.session_state.widget = ""
def get_result_text_es_pt (list_entity, text, lang):
result_words = []
tmp_word = ""
if lang == "es":
punc_tags = ['¿', '?', '¡', '!', ',', '.', ':']
else:
punc_tags = ['?', '!', ',', '.', ':']
for idx, entity in enumerate(list_entity):
tag = entity["entity"]
word = entity["word"]
start = entity["start"]
end = entity["end"]
# check punctuation
punc_in = next((p for p in punc_tags if p in tag), "")
subword = False
# check subwords
if word[0] == "#":
subword = True
if tmp_word == "":
p_s = list_entity[idx-1]["start"]
p_e = list_entity[idx-1]["end"]
tmp_word = text[p_s:p_e] + text[start:end]
else:
tmp_word = tmp_word + text[start:end]
word = tmp_word
else:
tmp_word = ""
word = text[start:end]
if tag == "l":
word = word
elif tag == "u":
word = word.capitalize()
# case with punctuation
else:
if tag[-1] == "l":
word = (punc_in + word) if punc_in in ["¿", "¡"] else (word + punc_in)
elif tag[-1] == "u":
word = (punc_in + word.capitalize()) if punc_in in ["¿", "¡"] else (word.capitalize() + punc_in)
if tag != "l":
word = '<span style="font-weight:bold; color:rgb(142, 208, 129);">' + word + '</span>'
if subword == True:
result_words[-1] = word
else:
result_words.append(word)
return " ".join(result_words)
def get_result_text_ca (list_entity, text):
result_words = []
punc_tags = ['?', '!', ',', '.', ':']
tmp_word = ""
for idx, entity in enumerate(list_entity):
start = entity["start"]
end = entity["end"]
tag = entity["entity"]
word = entity["word"]
# check punctuation
punc_in = next((p for p in punc_tags if p in tag), "")
subword = False
# check subwords
if word[0] != "Ġ":
subword = True
if tmp_word == "":
p_s = list_entity[idx-1]["start"]
p_e = list_entity[idx-1]["end"]
tmp_word = text[p_s:p_e] + text[start:end]
else:
tmp_word = tmp_word + text[start:end]
word = tmp_word
else:
tmp_word = ""
word = text[start:end]
if tag == "l":
word = word
elif tag == "u":
word = word.capitalize()
# case with punctuation
else:
if tag[-1] == "l":
word = (punc_in + word) if punc_in in ["¿", "¡"] else (word + punc_in)
elif tag[-1] == "u":
word = (punc_in + word.capitalize()) if punc_in in ["¿", "¡"] else (word.capitalize() + punc_in)
if tag != "l":
word = '<span style="font-weight:bold; color:rgb(142, 208, 129);">' + word + '</span>'
if subword == True:
result_words[-1] = word
else:
result_words.append(word)
return " ".join(result_words)
if __name__ == "__main__":
if "text" not in st.session_state:
st.session_state.text = ""
st.title('Sanivert Punctuation And Capitalization Restoration')
model_es = AutoModelForTokenClassification.from_pretrained("VOCALINLP/spanish_capitalization_punctuation_restoration_sanivert")
tokenizer_es = AutoTokenizer.from_pretrained("VOCALINLP/spanish_capitalization_punctuation_restoration_sanivert")
pipe_es = pipeline("token-classification", model=model_es, tokenizer=tokenizer_es)
model_ca = AutoModelForTokenClassification.from_pretrained("VOCALINLP/catalan_capitalization_punctuation_restoration_sanivert")
tokenizer_ca = AutoTokenizer.from_pretrained("VOCALINLP/catalan_capitalization_punctuation_restoration_sanivert")
pipe_ca = pipeline("token-classification", model=model_ca, tokenizer=tokenizer_ca)
model_pt = AutoModelForTokenClassification.from_pretrained("VOCALINLP/portuguese_capitalization_punctuation_restoration_sanivert")
tokenizer_pt = AutoTokenizer.from_pretrained("VOCALINLP/portuguese_capitalization_punctuation_restoration_sanivert")
pipe_pt = pipeline("token-classification", model=model_pt, tokenizer=tokenizer_pt)
st.sidebar.image("vocali_logo.jpg")
st.sidebar.subheader("Parque Científico de Murcia, Carretera de Madrid km 388. Complejo de Espinardo, 30100 Murcia")
input_text = st.selectbox(
label = "Choose an language",
options = ["Spanish", "Portuguese", "Catalan"]
)
st.subheader("Enter the text to be analyzed.")
st.text_input('Enter text', key='widget', on_change=clear_text) #text is stored in this variable
text = st.session_state.text
if input_text == "Spanish":
result_pipe = pipe_es(text)
out = get_result_text_es_pt(result_pipe, text, "es")
elif input_text == "Portuguese":
result_pipe = pipe_pt(text)
out = get_result_text_es_pt(result_pipe, text, "pt")
elif input_text == "Catalan":
result_pipe = pipe_ca(text)
out = get_result_text_ca(result_pipe, text)
st.markdown(out, unsafe_allow_html=True)
text = ""