Spaces:
Running
Running
| import random | |
| import re | |
| from poems import SAMPLE_POEMS | |
| import langid | |
| import numpy as np | |
| import streamlit as st | |
| import torch | |
| from icu_tokenizer import Tokenizer | |
| from transformers import pipeline | |
| MODELS = { | |
| "ALBERTI": "flax-community/alberti-bert-base-multilingual-cased", | |
| "mBERT": "bert-base-multilingual-cased" | |
| } | |
| TOPK = 50 | |
| st.set_page_config(layout="wide") | |
| def mask_line(line, language="es", restrictive=True): | |
| tokenizer = Tokenizer(lang=language) | |
| token_list = tokenizer.tokenize(line) | |
| if lang != "zh": | |
| restrictive = not all([len(token) <= 3 for token in token_list]) | |
| random_num = random.randint(0, len(token_list) - 1) | |
| random_word = token_list[random_num] | |
| if not restrictive: | |
| token_list[random_num] = "[MASK]" | |
| masked_l = " ".join(token_list) | |
| return masked_l | |
| elif len(random_word) > 3 or (lang == "zh" and random_word.isalpha()): | |
| token_list[random_num] = "[MASK]" | |
| masked_l = " ".join(token_list) | |
| return masked_l | |
| else: | |
| return mask_line(line, language) | |
| def filter_candidates(candidates, get_any_candidate=False): | |
| cand_list = [] | |
| score_list = [] | |
| for candidate in candidates: | |
| if not get_any_candidate and candidate["token_str"][:2] != "##" and candidate["token_str"].isalpha(): | |
| cand = candidate["sequence"] | |
| score = candidate["score"] | |
| cand_list.append(cand) | |
| score_list.append('{0:.5f}'.format(score)) | |
| elif get_any_candidate: | |
| cand = candidate["sequence"] | |
| score = candidate["score"] | |
| cand_list.append(cand) | |
| score_list.append('{0:.5f}'.format(score)) | |
| if len(score_list) == TOPK: | |
| break | |
| if len(cand_list) < 1: | |
| return filter_candidates(candidates, get_any_candidate=True) | |
| else: | |
| return cand_list[0] | |
| def infer_candidates(nlp, line): | |
| line = re.sub("–", "-", line) | |
| line = re.sub("—", "-", line) | |
| line = re.sub("’", "'", line) | |
| line = re.sub("…", "...", line) | |
| inputs = nlp._parse_and_tokenize(line) | |
| outputs = nlp._forward(inputs, return_tensors=True) | |
| input_ids = inputs["input_ids"][0] | |
| masked_index = torch.nonzero(input_ids == nlp.tokenizer.mask_token_id, | |
| as_tuple=False) | |
| logits = outputs[0, masked_index.item(), :] | |
| probs = logits.softmax(dim=0) | |
| values, predictions = probs.topk(TOPK) | |
| result = [] | |
| for v, p in zip(values.tolist(), predictions.tolist()): | |
| tokens = input_ids.numpy() | |
| tokens[masked_index] = p | |
| # Filter padding out: | |
| tokens = tokens[np.where(tokens != nlp.tokenizer.pad_token_id)] | |
| l = [] | |
| token_list = [nlp.tokenizer.decode([token], skip_special_tokens=True) for token in tokens] | |
| for idx, token in enumerate(token_list): | |
| if token.startswith('##'): | |
| l[-1] += token[2:] | |
| elif idx == masked_index.item(): | |
| l += ['<b style="color: #ff0000;">', token, "</b>"] | |
| else: | |
| l += [token] | |
| sequence = " ".join(l).strip() | |
| result.append( | |
| { | |
| "sequence": sequence, | |
| "score": v, | |
| "token": p, | |
| "token_str": nlp.tokenizer.decode(p), | |
| "masked_index": masked_index.item() | |
| } | |
| ) | |
| return result | |
| def rewrite_poem(poem, ml_model=MODELS["ALBERTI"], masking=True, language="es"): | |
| nlp = pipeline("fill-mask", model=ml_model) | |
| unmasked_lines = [] | |
| masked_lines = [] | |
| for line in poem: | |
| if line == "": | |
| unmasked_lines.append("") | |
| masked_lines.append("") | |
| continue | |
| if masking: | |
| masked_line = mask_line(line, language) | |
| else: | |
| masked_line = line | |
| masked_lines.append(masked_line) | |
| unmasked_line_candidates = infer_candidates(nlp, masked_line) | |
| unmasked_line = filter_candidates(unmasked_line_candidates) | |
| unmasked_lines.append(unmasked_line) | |
| unmasked_poem = "<br>".join(unmasked_lines) | |
| return unmasked_poem, masked_lines | |
| instructions_text_0 = st.sidebar.markdown( | |
| """# ALBERTI vs BERT 🥊 | |
| We present ALBERTI, our BERT-based multilingual model for poetry.""") | |
| instructions_text_1 = st.sidebar.markdown( | |
| """We have trained bert on a huge (for poetry, that is) corpus of | |
| multilingual poetry to try to get a more 'poetic' model. This is the result | |
| of our work. | |
| You can find more information on the [project's site](https://huggingface.co/flax-community/alberti-bert-base-multilingual-cased)""") | |
| sample_chooser = st.sidebar.selectbox( | |
| "Choose a poem", | |
| list(SAMPLE_POEMS.keys()) | |
| ) | |
| instructions_text_2 = st.sidebar.markdown("""# How to use | |
| You can choose from a list of example poems in Spanish, English, French, German, | |
| Chinese and Arabic, but you can also paste a poem, or write it yourself! | |
| Then click on 'Rewrite!' to do the masking and the fill-mask task on the chosen | |
| poem, randomly masking one word per verse, and get the two new versions for each of the models. | |
| The list of languages used on the training of ALBERTI are: | |
| * Arabic | |
| * Chinese | |
| * Czech | |
| * English | |
| * Finnish | |
| * French | |
| * German | |
| * Hungarian | |
| * Italian | |
| * Portuguese | |
| * Russian | |
| * Spanish""") | |
| col1, col2, col3 = st.columns(3) | |
| st.markdown( | |
| """ | |
| <style> | |
| label { | |
| font-size: 1rem !important; | |
| font-weight: bold !important; | |
| } | |
| .block-container { | |
| padding-left: 1rem !important; | |
| padding-right: 1rem !important; | |
| } | |
| </style> | |
| """, unsafe_allow_html=True) | |
| if sample_chooser: | |
| model_list = set(MODELS.values()) | |
| user_input = col1.text_area("Input poem", | |
| "\n".join(SAMPLE_POEMS[sample_chooser]), | |
| height=600) | |
| poem = user_input.split("\n") | |
| rewrite_button = col1.button("Rewrite!") | |
| if "[MASK]" in user_input or "<mask>" in user_input: | |
| col1.error("You don't have to mask the poem, we'll do it for you!") | |
| if rewrite_button: | |
| lang = langid.classify(user_input)[0] | |
| unmasked_poem, masked_poem = rewrite_poem(poem, language=lang) | |
| user_input_2 = col2.write(f"""<b>Output poem from ALBERTI</b> | |
| {unmasked_poem}""", unsafe_allow_html=True) | |
| unmasked_poem_2, _ = rewrite_poem(masked_poem, ml_model=MODELS["mBERT"], | |
| masking=False) | |
| user_input_3 = col3.write(f"""<b>Output poem from mBERT</b> | |
| {unmasked_poem_2}""", unsafe_allow_html=True) | |