twigs/simplifier

1import streamlit as st2from transformers import AutoTokenizer, AutoModelForSequenceClassification, BartTokenizer, BartForConditionalGeneration, pipeline3import numpy as np4import torch5import re6from textstat import textstat7 8 9MAX_LEN = 25610NUM_BEAMS = 411EARLY_STOPPING = True12N_OUT = 413 14 15cwi_tok = AutoTokenizer.from_pretrained('twigs/cwi-regressor')16cwi_model = AutoModelForSequenceClassification.from_pretrained(17    'twigs/cwi-regressor')18simpl_tok = BartTokenizer.from_pretrained('twigs/bart-text2text-simplifier')19simpl_model = BartForConditionalGeneration.from_pretrained(20    'twigs/bart-text2text-simplifier')21cwi_pipe = pipeline('text-classification', model=cwi_model,22                    tokenizer=cwi_tok, function_to_apply='none')23fill_pipe = pipeline('fill-mask', top_k=1)24 25 26def id_replace_complex(s, threshold=0.2):27 28  # get all tokens29  tokens = re.compile('\w+').findall(s)30  cands = [f"{t}. {s}" for t in tokens]31  # get complex tokens32  # if score >= threshold select tokens[idx]33  compl_tok = [tokens[idx] for idx, x in enumerate(34      cwi_pipe(cands)) if x['score'] >= threshold]35  36  masked = [s[:s.index(t)] + '<mask>' + s[s.index(t)+len(t):] for t in compl_tok]37  cands = fill_pipe(masked)38  # structure is different in 1 vs n complex words39  replacements = [el['token_str'] if type(40      el) == dict else el[0]['token_str'] for el in cands]41  # some tokens get prefixed with space42  replacements = [tok if tok.find(' ') == -1 else tok[1:]43                  for tok in replacements]44 45  for i, el in enumerate(compl_tok):46    idx = s.index(el)47    s = s[:idx] + replacements[i] + s[idx+len(el):]48  49  return s, compl_tok, replacements50 51def generate_candidate_text(s, model, tokenizer, tokenized=False):52 53 54  out = simpl_tok([s], max_length=256, padding="max_length",  truncation=True,55                  return_tensors='pt') if not tokenized else s56 57  generated_ids = model.generate(58      input_ids=out['input_ids'],59      attention_mask=out['attention_mask'],60      use_cache=True,61      decoder_start_token_id=simpl_model.config.pad_token_id,62      num_beams=NUM_BEAMS,63      max_length=MAX_LEN,64      early_stopping=EARLY_STOPPING,65      num_return_sequences=N_OUT66  )67 68  return [tokenizer.decode(ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)[69      1:] for ids in generated_ids]70 71 72def rank_candidate_text(sentences):73    fkgl_scores = [textstat.flesch_kincaid_grade(s) for s in sentences]74    return sentences[np.argmin(fkgl_scores)]75 76 77def full_pipeline(source, simpl_model, simpl_tok, tokens, lexical=False):78  79  modified, complex_words, replacements  = id_replace_complex(source, threshold=0.2) if lexical else (source, None, None)80  cands = generate_candidate_text(tokens+modified, simpl_model, simpl_tok)81  output = rank_candidate_text(cands)82  return output, complex_words, replacements83  84def main():85 86    aug_tok = ['c_', 'lev_', 'dep_', 'rank_', 'rat_', 'n_syl_']87    base_tokens = ['CharRatio', 'LevSim', 'DependencyTreeDepth',88            'WordComplexity', 'WordRatio', 'NumberOfSyllables']89 90    default_values = [0.8, 0.6, 0.9, 0.8, 0.9, 1.9]91    user_values = default_values92    tok_values = dict((t, default_values[idx]) for idx, t in enumerate(base_tokens))93 94    example_sentences = ["A matchbook is a small cardboard folder (matchcover) enclosing a quantity of matches and having a coarse striking surface on the exterior.",95                        "If there are no strong land use controls, buildings are built along a bypass, converting it into an ordinary town road, and the bypass may eventually become as congested as the local streets it was intended to avoid.",96                        "Plot Captain Caleb Holt (Kirk Cameron) is a firefighter in Albany, Georgia and firmly keeps the cardinal rule of all firemen, \"Never leave your partner behind\".",97                        "Britpop emerged from the British independent music scene of the early 1990s and was characterised by bands influenced by British guitar pop music of the 1960s and 1970s."]98 99 100    st.title("Make it Simple")101 102    with st.expander("Example sentences"):103        for s in example_sentences:104            st.code(body=s)105 106 107    with st.form(key="simplify"):108        input_sentence = st.text_area("Original sentence")109    110        lexical = st.checkbox("Identify and replace complex words", value=True)111 112        tok = st.multiselect(113            label="Tokens to augment the sentence", options=base_tokens, default=base_tokens)114        if (tok):115            st.text("Select the desired intensity")116            for idx, t in enumerate(tok):117                user_values[idx] = st.slider(118                    t, min_value=0., max_value=1., value=tok_values[t], step=0.1, key=t)119 120        submit = st.form_submit_button("Process")121        if (submit):122            123            tokens = " ".join([t+str(v) for t, v in zip(aug_tok, user_values)]) + " "124            output, words, replacements = full_pipeline(input_sentence, simpl_model, simpl_tok, tokens, lexical)125            126    127            c1, c2, c3 = st.columns([1,1,2])128 129            with c1:130                st.markdown("#### Words identified as complex")131                if words:132                    for w in words:133                        st.markdown(f"* {w}")134 135                else:136                    st.markdown("None :smile:")137 138            with c2:139                st.markdown("#### Their mask-predicted replacement")140                if replacements:141                    for w in replacements:142                        st.markdown(f"* {w}")143 144                else:145                    st.markdown("None :smile:")146 147            with c3:148                st.markdown(f"#### Original Sentence:\n > {input_sentence}") 149                st.markdown(f"#### Output Sentence:\n > {output}") 150 151 152if __name__ == '__main__':153    main()154