critic/critic.py

import sys
import torch
import random
import hashlib
import numpy as np
from tqdm import tqdm
from transformers import GPT2Tokenizer, GPT2Model, GPT2LMHeadModel
from transformers import OPTForCausalLM, GPTNeoForCausalLM
from transformers import RobertaTokenizer, RobertaForCausalLM, RobertaConfig
from transformers import XLMRobertaTokenizer, XLMRobertaForCausalLM, XLMRobertaConfig
from transformers import BartTokenizer, BartForCausalLM
import nltk
import pandas as pd
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('Enter any sentence in the text box, press submit, and see the grammatical scoring and judgement results outputted by LM-Critic using different LMs displayed below.')

def get_gpt2_loss(model, tokenizer, input_ids, attention_mask, labels):
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)
        lm_logits = outputs[1] #[bsize, seqlen, vocab]
        if labels is not None:
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            shift_mask = attention_mask[..., 1:].contiguous()
            loss_fct = torch.nn.CrossEntropyLoss(reduction='none')
            bsize, seqlen = input_ids.size()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)).view(bsize, seqlen-1)
            loss = (loss * shift_mask).sum(dim=1) #[bsize, ]
        return loss


MAX_LENGTH = 66

def run_gpt2(sents, model, tokenizer, cuda=False, model_name=None):
    assert isinstance(sents, list)
    _sents = [tokenizer.bos_token + s for s in sents]
    inputs = tokenizer(_sents, return_tensors="pt", padding=True)
    if inputs['input_ids'].size(1) > MAX_LENGTH:
        return None
    if cuda:
        inputs = {k: v.cuda() for k, v in inputs.items()}
    loss = get_gpt2_loss(model, tokenizer, input_ids=inputs['input_ids'], attention_mask=inputs['attention_mask'], labels=inputs['input_ids'])
    logps = - loss.detach().cpu()
    return logps


def gpt2_critic_char_level_only(sent, verbose=1, cuda=False, fp16=True, seed='auto', n_samples=100):
    return_string = []
    if seed == 'auto':
        seed = int(hashlib.md5(sent.encode()).hexdigest(), 16) % (2**32) #Seed must be between 0 and 2**32 - 1
    if verbose > 1:
        print ('seed', seed)
    np.random.seed(seed); random.seed(seed)
    is_good = True
    for _ in range(1):
        sent_perturbations = get_local_neighbors_char_level(sent, max_n_samples=n_samples)
        if verbose > 1:
            print ("#sent_perturbations (char-level)", len(sent_perturbations))
            return_string.append(f"#sent_perturbations (char-level){len(sent_perturbations)}\n")
        sents = [sent] + list(sent_perturbations)
        if fp16:
            with torch.cuda.amp.autocast():
                logps = run_gpt2(sents, cuda)
        else:
            logps = run_gpt2(sents, cuda)
        if logps is None:
            if verbose:
                print ('Invalid input. Maybe the sentence is too long.')
                return_string.append('Invalid input. Maybe the sentence is too long.\n')
            return None
        best_idx = int(logps.argmax())
        if best_idx != 0:
            is_good = False
            break
    if verbose:
        if is_good:
            print ('Good! Your sentence log(p) = {:.3f}'.format(float(logps[0])))
            return_string.append('Good! Your sentence log(p) = {:.3f}\n'.format(float(logps[0])))
        else:
            print ('Bad! Your sentence log(p) = {:.3f}'.format(float(logps[0])))
            return_string.append('Bad! Your sentence log(p) = {:.3f}\n'.format(float(logps[0])))
            print ('Neighbor sentence with highest log(p): {} (= {:.3f})'.format(sents[best_idx], float(logps[best_idx])))
            return_string.append('Neighbor sentence with highest log(p): {} (= {:.3f})\n'.format(sents[best_idx], float(logps[best_idx])))
    counter_example = None
    if not is_good:
        counter_example = [sents[best_idx], float(logps[best_idx])]
    return is_good, float(logps[0]), counter_example


def gpt2_critic(sent, model, tokenizer, verbose=1, cuda=False, fp16=True, seed='auto', n_samples=100, word_level_mode='refine'):
    return_string = []
    if seed == 'auto':
        seed = int(hashlib.md5(sent.encode()).hexdigest(), 16) % (2**32) #Seed must be between 0 and 2**32 - 1
    if verbose > 1:
        print ('seed', seed)
        return_string.append(f'seed{seed}\n')
    np.random.seed(seed); random.seed(seed)
    sent_toked = spacy_tokenize_gec(sent)
    is_good = True
    for _ in range(1):
        sent_perturbations_w, orig_sent = get_local_neighbors_word_level(sent_toked, max_n_samples=n_samples//2, mode=word_level_mode)
        sent_perturbations_c = get_local_neighbors_char_level(orig_sent, max_n_samples=n_samples//2)
        if verbose > 1:
            print ("#sent_perturbations (char-level)", len(sent_perturbations_c))
            return_string.append("#sent_perturbations (char-level)\n", len(sent_perturbations_c))
            print ("#sent_perturbations (word-level)", len(sent_perturbations_w))
            return_string.append("#sent_perturbations (word-level)\n", len(sent_perturbations_w))
        sents = [orig_sent] + list(sent_perturbations_c.union(sent_perturbations_w))
        if fp16:
            with torch.cuda.amp.autocast():
                logps = run_gpt2(sents, model, tokenizer, cuda)
        else:
            logps = run_gpt2(sents, model, tokenizer, cuda)
        if logps is None:
            if verbose:
                print ('Invalid input. Maybe the sentence is too long.')
                return_string.append('Invalid input. Maybe the sentence is too long.\n')
            return None
        best_idx = int(logps.argmax())
        if best_idx != 0:
            is_good = False
            break
    if verbose:
        if is_good:
            print ('Good! Your sentence log(p) = {:.3f}'.format(float(logps[0])))
            return_string.append('Good! Your sentence log(p) = {:.3f}\n'.format(float(logps[0])))
        else:
            print ('Bad! Your sentence log(p) = {:.3f}'.format(float(logps[0])))
            return_string.append('Bad! Your sentence log(p) = {:.3f}\n'.format(float(logps[0])))
            print ('Neighbor sentence with highest log(p): {} (= {:.3f})'.format(sents[best_idx], float(logps[best_idx])))
            return_string.append('Neighbor sentence with highest log(p): {} (= {:.3f})\n'.format(sents[best_idx], float(logps[best_idx])))
    counter_example = None
    if not is_good:
        counter_example = [sents[best_idx], float(logps[best_idx])]
    return is_good, float(logps[0]), counter_example, return_string

def gpt2():
    ## GPT-2 LM (original LM-critic)
    placeholder_lm_name = st.empty()
    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()
    placeholder_lm_name.empty()
    st.session_state["model_gpt2"] = model_gpt2
    st.session_state["tokenizer_gpt2"] = tokenizer_gpt2
    st.session_state["nice_name_gpt2"] = nice_name_gpt2

def opt():
    ## OPT LM
    placeholder_lm_name = st.empty()
    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()
    placeholder_lm_name.empty()
    st.session_state["model_opt"] = model_opt
    st.session_state["tokenizer_opt"] = tokenizer_opt
    st.session_state["nice_name_opt"] = nice_name_opt

def gpt_neo():
    ## GPT NEO
    placeholder_lm_name = st.empty()
    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()
    placeholder_lm_name.empty()
    st.session_state["model_gptneo"] = model_gptneo
    st.session_state["tokenizer_gptneo"] = tokenizer_gptneo
    st.session_state["nice_name_gptneo"] = nice_name_gptneo

def roberta():
    ## RoBERTa
    placeholder_lm_name = st.empty()
    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()
    placeholder_lm_name.empty()
    st.session_state["model_roberta"] = model_roberta
    st.session_state["tokenizer_roberta"] = tokenizer_roberta
    st.session_state["nice_name_roberta"] = nice_name_roberta

def bart():
    ## BART
    placeholder_lm_name = st.empty()
    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()
    placeholder_lm_name.empty()
    st.session_state["model_bart"] = model_bart
    st.session_state["tokenizer_bart"] = tokenizer_bart
    st.session_state["nice_name_bart"] = nice_name_bart

def xlm_roberta():
    ## XLM RoBERTa
    placeholder_lm_name = st.empty()
    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()
    placeholder_lm_name.empty()
    st.session_state["model_xlmroberta"] = model_xlmroberta
    st.session_state["tokenizer_xlmroberta"] = tokenizer_xlmroberta
    st.session_state["nice_name_xlmroberta"] = nice_name_xlmroberta

def main():
    form = st.form(key='my_form')
    sent = form.text_input(label='Enter a sentence:', value="")
    submit = form.form_submit_button(label='Submit')

    if submit and sent != '':
        st.markdown(f"**Input Sentence**: {sent}")
        results = {}

        with st.spinner('Running with GPT-2 LM...'):
            ## GPT-2 LM (original LM-critic)
            if "nice_name_gpt2" not in st.session_state:
                gpt2()
            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[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
            if "nice_name_opt" not in st.session_state:
                opt()
            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[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
            if "nice_name_gptneo" not in st.session_state:
                gpt_neo()
            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[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
            if "nice_name_roberta" not in st.session_state:
                roberta()
            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[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...'):
            ## BART
            if "nice_name_bart" not in st.session_state:
                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[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
            if "nice_name_xlmroberta" not in st.session_state:
                xlm_roberta()
            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[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',
            columns=['Judgement', 'Score (log(p))', 'Neighbor sentence with highest score (log(p))', 'Neighbor sentence score (log(p))'])
        st.markdown("**Tabular summary of results:**")
        st.table(df)

        st.write("Input another sentence!")

if __name__ == '__main__':
    main()