Fix runtime memory issues by caching and using session state

critic/critic.py

@@ -16,6 +16,11 @@ nltk.download('punkt')
 sys.path.insert(0, '.')
 from critic.perturbations import get_local_neighbors_char_level, get_local_neighbors_word_level
 from utils.spacy_tokenizer import spacy_tokenize_gec
+import streamlit as st
+
+st.subheader('Exploring Unsupervised Grammatical Error Correction with Transformer-Based Models')
+st.write('This live demonstration is adapted from the paper [LM-Critic: Language Models for Unsupervised Grammatical Error Correction](https://aclanthology.org/2021.emnlp-main.611.pdf) (EMNLP 2021) by Michihiro Yasunaga, Jure Leskovec, Percy Liang.')
+st.write('The below demo first loads several LMs that we use in the LM-Critic. You will be prompted to enter a sentence which will then be scored by each of the LM-Critics using different LMs.')
 
 def get_gpt2_loss(model, tokenizer, input_ids, attention_mask, labels):
     with torch.no_grad():
@@ -137,10 +142,117 @@ def gpt2_critic(sent, model, tokenizer, verbose=1, cuda=False, fp16=True, seed='
         counter_example = [sents[best_idx], float(logps[best_idx])]
     return is_good, float(logps[0]), counter_example, return_string
 
+@st.cache(suppress_st_warning=True)
+def init_lms():
+    placeholder_lm_name = st.empty()
+    my_bar = st.progress(10)
+    prog = 10
+
+    ## GPT-2 LM (original LM-critic)
+    model_name_gpt2 = 'gpt2'
+    nice_name_gpt2 = "GPT-2"
+    placeholder_lm_name.text(f"Initializing {nice_name_gpt2}...")
+    tokenizer_gpt2 = GPT2Tokenizer.from_pretrained(model_name_gpt2)
+    tokenizer_gpt2.pad_token = tokenizer_gpt2.eos_token
+    model_gpt2 = GPT2LMHeadModel.from_pretrained(model_name_gpt2)
+    model_gpt2.eval()
+    model_gpt2.cpu()
+    st.session_state["model_gpt2"] = model_gpt2
+    st.session_state["tokenizer_gpt2"] = tokenizer_gpt2
+    st.session_state["nice_name_gpt2"] = nice_name_gpt2
+
+    prog += 10
+    my_bar.progress(prog)    
+
+    ## OPT LM
+    model_name_opt = "facebook/opt-350m"
+    nice_name_opt = "OPT"
+    placeholder_lm_name.text(f"Initializing {nice_name_opt}...")
+    model_opt = OPTForCausalLM.from_pretrained(model_name_opt)
+    tokenizer_opt = GPT2Tokenizer.from_pretrained(model_name_opt)
+    tokenizer_opt.pad_token = tokenizer_opt.eos_token
+    model_opt.eval()
+    model_opt.cpu()
+    st.session_state["model_opt"] = model_opt
+    st.session_state["tokenizer_opt"] = tokenizer_opt
+    st.session_state["nice_name_opt"] = nice_name_opt
+
+    prog += 10
+    my_bar.progress(prog)  
+
+    ## GPT NEO
+    model_name_gptneo = "EleutherAI/gpt-neo-1.3B"
+    nice_name_gptneo = "GPT NEO"
+    placeholder_lm_name.text(f"Initializing {nice_name_gptneo}...")
+    model_gptneo = GPTNeoForCausalLM.from_pretrained(model_name_gptneo)
+    tokenizer_gptneo = GPT2Tokenizer.from_pretrained(model_name_gptneo)
+    tokenizer_gptneo.pad_token = tokenizer_gptneo.eos_token
+    model_gptneo.eval()
+    model_gptneo.cpu()
+    st.session_state["model_gptneo"] = model_gptneo
+    st.session_state["tokenizer_gptneo"] = tokenizer_gptneo
+    st.session_state["nice_name_gptneo"] = nice_name_gptneo
+
+    prog += 10
+    my_bar.progress(prog)  
+
+    ## RoBERTa
+    model_name_roberta = "roberta-base"
+    nice_name_roberta = "RoBERTa"
+    placeholder_lm_name.text(f"Initializing {nice_name_roberta}...")
+    tokenizer_roberta = RobertaTokenizer.from_pretrained(model_name_roberta)
+    config_roberta = RobertaConfig.from_pretrained(model_name_roberta)
+    config_roberta.is_decoder = True
+    model_roberta = RobertaForCausalLM.from_pretrained(model_name_roberta, config=config_roberta)
+    tokenizer_roberta.pad_token = tokenizer_roberta.eos_token
+    model_roberta.eval()
+    model_roberta.cpu()
+    st.session_state["model_roberta"] = model_gptneo
+    st.session_state["tokenizer_roberta"] = tokenizer_roberta
+    st.session_state["nice_name_roberta"] = nice_name_roberta
+
+    prog += 10
+    my_bar.progress(prog)  
+
+    ## BART
+    model_name_bart = "facebook/bart-base"
+    nice_name_bart = "BART"
+    placeholder_lm_name.text(f"Initializing {nice_name_bart}...")
+    tokenizer_bart = BartTokenizer.from_pretrained(model_name_bart)
+    model_bart = BartForCausalLM.from_pretrained(model_name_bart, add_cross_attention=False)
+    assert model_bart.config.is_decoder, f"{model_bart.__class__} has to be configured as a decoder."
+    tokenizer_bart.pad_token = tokenizer_bart.eos_token
+    model_bart.eval()
+    model_bart.cpu()
+    st.session_state["model_bart"] = model_bart
+    st.session_state["tokenizer_bart"] = tokenizer_bart
+    st.session_state["nice_name_bart"] = nice_name_bart
+
+    prog += 10
+    my_bar.progress(prog)  
+
+    ## XLM RoBERTa
+    model_name_xlmroberta = 'xlm-roberta-base'
+    nice_name_xlmroberta = 'XLM RoBERTa'
+    placeholder_lm_name.text(f"Initializing {nice_name_xlmroberta}...")
+    tokenizer_xlmroberta = XLMRobertaTokenizer.from_pretrained(model_name_xlmroberta)
+    config_xlmroberta = XLMRobertaConfig.from_pretrained(model_name_xlmroberta)
+    config_xlmroberta.is_decoder = True
+    model_xlmroberta = XLMRobertaForCausalLM.from_pretrained(model_name_xlmroberta, config=config_xlmroberta)
+    tokenizer_xlmroberta.pad_token = tokenizer_xlmroberta.eos_token
+    model_xlmroberta.eval()
+    model_xlmroberta.cpu()
+    st.session_state["model_xlmroberta"] = model_xlmroberta
+    st.session_state["tokenizer_xlmroberta"] = tokenizer_xlmroberta
+    st.session_state["nice_name_xlmroberta"] = nice_name_xlmroberta
+
+    prog += 10
+    my_bar.progress(prog)
+    placeholder_lm_name.empty()
+    my_bar.empty()
 
 def main():
-    import streamlit as st
-    st.subheader('Exploring Unsupervised Grammatical Error Correction with Transformer-Based Models')
+    init_lms()
     sent = st.text_input('Enter a sentence:', value="")
 
     ### LMs we are trying:
@@ -150,92 +262,45 @@ def main():
 
         with st.spinner('Running with GPT-2 LM...'):
             ## GPT-2 LM (original LM-critic)
-            model_name = 'gpt2'
-            nice_name = "GPT-2"
-            tokenizer = GPT2Tokenizer.from_pretrained(model_name)
-            tokenizer.pad_token = tokenizer.eos_token
-            model = GPT2LMHeadModel.from_pretrained(model_name)
-            model.eval()
-            model.cpu()
-            is_good, score, counter_example, return_string_GPT2 = gpt2_critic(sent, model, tokenizer)
+            is_good, score, counter_example, return_string_GPT2 = gpt2_critic(sent, st.session_state['model_gpt2'], st.session_state['tokenizer_gpt2'])
         st.markdown("**Results with GPT-2 LM:**")
         st.write('\n'.join(return_string_GPT2))
-        results[nice_name] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
+        results[st.session_state['nice_name_gpt2']] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
 
         with st.spinner('Running with OPT LM...'):
             ## OPT LM
-            model_name = "facebook/opt-350m"
-            nice_name = "OPT"
-            model = OPTForCausalLM.from_pretrained("facebook/opt-350m")
-            tokenizer = GPT2Tokenizer.from_pretrained("facebook/opt-350m")
-            tokenizer.pad_token = tokenizer.eos_token
-            model.eval()
-            model.cpu()
-            is_good, score, counter_example, return_string_OPT = gpt2_critic(sent, model, tokenizer)
+            is_good, score, counter_example, return_string_OPT = gpt2_critic(sent, st.session_state['model_opt'], st.session_state['tokenizer_opt'])
         st.markdown("**Results with OPT LM:**")
         st.write('\n'.join(return_string_OPT))
-        results[nice_name] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
+        results[st.session_state['nice_name_opt']] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
 
         with st.spinner('Running with GPT NEO LM...'):
             ## GPT NEO
-            model_name = "EleutherAI/gpt-neo-1.3B"
-            nice_name = "GPT NEO"
-            model = GPTNeoForCausalLM.from_pretrained("EleutherAI/gpt-neo-1.3B")
-            tokenizer = GPT2Tokenizer.from_pretrained("EleutherAI/gpt-neo-1.3B")
-            tokenizer.pad_token = tokenizer.eos_token
-            model.eval()
-            model.cpu()
-            is_good, score, counter_example, return_string_GPTNEO = gpt2_critic(sent, model, tokenizer)
+            is_good, score, counter_example, return_string_GPTNEO = gpt2_critic(sent, st.session_state['model_gptneo'], st.session_state['tokenizer_gptneo'])
         st.markdown("**Results with GPT NEO LM:**")
         st.write('\n'.join(return_string_GPTNEO))
-        results[nice_name] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
+        results[st.session_state['nice_name_gptneo']] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
 
         with st.spinner('Running with RoBERTa LM...'):
             ## RoBERTa
-            model_name = "roberta-base"
-            nice_name = "RoBERTa"
-            tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
-            config = RobertaConfig.from_pretrained("roberta-base")
-            config.is_decoder = True
-            model = RobertaForCausalLM.from_pretrained("roberta-base", config=config)
-            tokenizer.pad_token = tokenizer.eos_token
-            model.eval()
-            model.cpu()
-            is_good, score, counter_example, return_string_RoBERTa = gpt2_critic(sent, model, tokenizer)
+            is_good, score, counter_example, return_string_RoBERTa = gpt2_critic(sent, st.session_state['model_roberta'], st.session_state['tokenizer_roberta'])
         st.markdown("**Results with RoBERTa LM:**")
         st.write('\n'.join(return_string_RoBERTa))
-        results[nice_name] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
+        results[st.session_state['nice_name_roberta']] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
 
         with st.spinner('Running with BART LM...'):
-            ## RoBERTa
-            model_name = "facebook/bart-base"
-            nice_name = "BART"
-            tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
-            model = BartForCausalLM.from_pretrained("facebook/bart-base", add_cross_attention=False)
-            assert model.config.is_decoder, f"{model.__class__} has to be configured as a decoder."
-            tokenizer.pad_token = tokenizer.eos_token
-            model.eval()
-            model.cpu()
-            is_good, score, counter_example, return_string_BART = gpt2_critic(sent, model, tokenizer)
+            ## BART
+            is_good, score, counter_example, return_string_BART = gpt2_critic(sent, st.session_state['model_bart'], st.session_state['tokenizer_bart'])
         st.markdown("**Results with BART LM:**")
         st.write('\n'.join(return_string_BART))
-        results[nice_name] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
+        results[st.session_state['nice_name_bart']] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
 
         with st.spinner('Running with XLM RoBERTa LM...'):
             ## XLM RoBERTa
-            model_name = 'xlm-roberta-base'
-            nice_name = 'XLM RoBERTa'
-            tokenizer = XLMRobertaTokenizer.from_pretrained("xlm-roberta-base")
-            config = XLMRobertaConfig.from_pretrained("xlm-roberta-base")
-            config.is_decoder = True
-            model = XLMRobertaForCausalLM.from_pretrained("xlm-roberta-base", config=config)
-            tokenizer.pad_token = tokenizer.eos_token
-            model.eval()
-            model.cpu()
-            is_good, score, counter_example, return_string_XLMRoBERTa = gpt2_critic(sent, model, tokenizer)
+            is_good, score, counter_example, return_string_XLMRoBERTa = gpt2_critic(sent, st.session_state['model_xlmroberta'], st.session_state['tokenizer_xlmroberta'])
         st.markdown("**Results with XLM RoBERTa LM:**")
         st.write('\n'.join(return_string_XLMRoBERTa))
-        results[nice_name] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
+        results[st.session_state['nice_name_xlmroberta']] = ["Good" if is_good else "Bad", str(round(score, 3)), "N/A" if not counter_example else str(counter_example[0]), "N/A" if not counter_example else str(round(counter_example[1], 3))]
 
         df = pd.DataFrame.from_dict(results, 
             orient = 'index',
@@ -243,7 +308,7 @@ def main():
         st.markdown("**Tabular summary of results:**")
         st.table(df)
 
-        st.write("Done.")
+        st.write("Input another sentence!")
 
 if __name__ == '__main__':
     main()