Upload code with streamlit addition

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Michihiro Yasunaga
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -5,7 +5,7 @@ colorFrom: pink
 colorTo: gray
 sdk: streamlit
 sdk_version: 1.2.0
-app_file: app.py
+app_file: critic/critic.py
 pinned: false
 license: afl-3.0
 ---

critic/PIE/word_level_perturb.py

@@ -0,0 +1,237 @@
+"""
+Word-level perturbation generator.
+
+Originally by https://github.com/awasthiabhijeet/PIE/tree/master/errorify
+"""
+import os
+import math
+import pickle
+import random
+import editdistance
+from numpy.random import choice as npchoice
+from collections import defaultdict
+
+
+try:
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+except:
+    dir_path = '.'
+
+VERBS = pickle.load(open(f'{dir_path}/verbs.p', 'rb'))
+COMMON_INSERTS = set(pickle.load(open(f'{dir_path}/common_inserts.p', 'rb'))) #common inserts *to fix a sent*
+COMMON_DELETES = pickle.load(open(f'{dir_path}/common_deletes.p','rb')) #common deletes *to fix a sent*
+_COMMON_REPLACES = pickle.load(open(f'{dir_path}/common_replaces.p', 'rb')) #common replacements *to errorify a sent*
+
+
+
+COMMON_REPLACES = {}
+for src in _COMMON_REPLACES:
+    for tgt in _COMMON_REPLACES[src]:
+        if (src=="'re" and tgt=="are") or (tgt=="'re" and src=="are"):
+            continue
+        ED = editdistance.eval(tgt, src)
+        if ED > 2:
+            continue
+        longer = max(len(src), len(tgt))
+        if float(ED)/longer >= 0.5:
+            continue
+        if tgt not in COMMON_REPLACES:
+            COMMON_REPLACES[tgt] = {}
+        COMMON_REPLACES[tgt][src] = _COMMON_REPLACES[src][tgt]
+
+
+VERBS_refine = defaultdict(list)
+for src in VERBS:
+    for tgt in VERBS[src]:
+        ED = editdistance.eval(tgt, src)
+        if ED > 2:
+            continue
+        longer = max(len(src), len(tgt))
+        if float(ED)/longer >= 0.5:
+            continue
+        VERBS_refine[src].append(tgt)
+
+
+
+class WordLevelPerturber_all:
+    def __init__(self, sentence: str):
+        self.original_sentence = sentence.rstrip()
+        self.sentence = self.original_sentence
+        self.tokenized = None
+        self.tokenize()
+
+    def tokenize(self):
+        self.tokenized = self.sentence.split()
+
+    def orig(self):
+        return self.original_sentence
+
+    def _insert(self):
+        """Insert a commonly deleted word."""
+        if len(self.tokenized) > 0:
+            insertable = list(range(len(self.tokenized)))
+            index = random.choice(insertable)
+            plist = list(COMMON_DELETES.values())
+            plistsum = sum(plist)
+            plist = [x / plistsum for x in plist]
+            # Choose a word
+            ins_word = npchoice(list(COMMON_DELETES.keys()), p=plist)
+            self.tokenized.insert(index,ins_word)
+        return ' '.join(self.tokenized)
+
+    def _mod_verb(self, redir=True):
+        if len(self.tokenized) > 0:
+            verbs = [i for i, w in enumerate(self.tokenized) if w in VERBS]
+            if not verbs:
+                if redir:
+                    return self._replace(redir=False)
+                return self.sentence
+            index = random.choice(verbs)
+            word = self.tokenized[index]
+            if not VERBS[word]:
+                return self.sentence
+            repl = random.choice(VERBS[word])
+            self.tokenized[index] = repl
+        return ' '.join(self.tokenized)
+
+    def _delete(self):
+        """Delete a commonly inserted word."""
+        if len(self.tokenized) > 1:
+            toks_len = len(self.tokenized)
+            toks = self.tokenized
+            deletable = [i for i, w in enumerate(toks) if w in COMMON_INSERTS]
+            if not deletable:
+                return self.sentence
+            index = random.choice(deletable)
+            del self.tokenized[index]
+        return ' '.join(self.tokenized)
+
+    def _replace(self, redir=True):
+        if len(self.tokenized) > 0:
+            deletable = [i for i, w in enumerate(self.tokenized) if (w in COMMON_REPLACES)]
+            if not deletable:
+                if redir:
+                    return self._mod_verb(redir=False)
+                return self.sentence
+            index = random.choice(deletable)
+            word = self.tokenized[index]
+            if not COMMON_REPLACES[word]:
+                return self.sentence
+            # Normalize probabilities
+            plist = list(COMMON_REPLACES[word].values())
+            plistsum = sum(plist)
+            plist = [x / plistsum for x in plist]
+            # Choose a word
+            repl = npchoice(list(COMMON_REPLACES[word].keys()), p=plist)
+            self.tokenized[index] = repl
+        return ' '.join(self.tokenized)
+
+    def perturb(self):
+        count = 1
+        orig_sent = self.sentence
+        for x in range(count):
+            perturb_probs = [.30,.30,.30,.10]
+            perturb_fun = npchoice([self._insert, self._mod_verb, self._replace, self._delete],p=perturb_probs)
+            self.sentence = perturb_fun()
+            self.tokenize()
+        res_sentence = self.sentence
+        self.sentence = self.original_sentence
+        self.tokenize()
+        return res_sentence
+
+
+class WordLevelPerturber_refine:
+    def __init__(self, sentence: str):
+        self.original_sentence = sentence.rstrip()
+        self.sentence = self.original_sentence
+        self.tokenized = None
+        self.tokenize()
+
+    def tokenize(self):
+        self.tokenized = self.sentence.split()
+
+    def orig(self):
+        return self.original_sentence
+
+    def _insert(self):
+        """Insert a commonly deleted word."""
+        if len(self.tokenized) > 0:
+            insertable = list(range(len(self.tokenized)))
+            index = random.choice(insertable)
+            plist = list(COMMON_DELETES.values())
+            plistsum = sum(plist)
+            plist = [x / plistsum for x in plist]
+            # Choose a word
+            ins_word = npchoice(list(COMMON_DELETES.keys()), p=plist)
+            self.tokenized.insert(index,ins_word)
+        return ' '.join(self.tokenized)
+
+    def _mod_verb(self, redir=True):
+        if len(self.tokenized) > 0:
+            verbs = [i for i, w in enumerate(self.tokenized) if w in VERBS_refine]
+            if not verbs:
+                if redir:
+                    return self._replace(redir=False)
+                return self.sentence
+            index = random.choice(verbs)
+            word = self.tokenized[index]
+            if not VERBS_refine[word]:
+                return self.sentence
+            repl = random.choice(VERBS_refine[word])
+            self.tokenized[index] = repl
+
+        return ' '.join(self.tokenized)
+
+    def _delete(self):
+        """Delete a commonly inserted word."""
+        if len(self.tokenized) > 1:
+            toks_len = len(self.tokenized)
+            toks = self.tokenized
+            deletable = [i for i, w in enumerate(toks) if (w in COMMON_INSERTS) and (i>0 and toks[i-1].lower() == toks[i].lower())]
+            if not deletable:
+                return self.sentence
+            index = random.choice(deletable)
+            del self.tokenized[index]
+        return ' '.join(self.tokenized)
+
+    def _replace(self, redir=True):
+        def _keep(i,w):
+            if w.lower() in {"not", "n't"}:
+                return True
+            return False
+
+        if len(self.tokenized) > 0:
+            deletable = [i for i, w in enumerate(self.tokenized) if (w in COMMON_REPLACES) and (not _keep(i,w))]
+            if not deletable:
+                if redir:
+                    return self._mod_verb(redir=False)
+                return self.sentence
+
+            index = random.choice(deletable)
+            word = self.tokenized[index]
+            if not COMMON_REPLACES[word]:
+                return self.sentence
+
+            # Normalize probabilities
+            plist = list(COMMON_REPLACES[word].values())
+            plistsum = sum(plist)
+            plist = [x / plistsum for x in plist]
+
+            # Choose a word
+            repl = npchoice(list(COMMON_REPLACES[word].keys()), p=plist)
+            self.tokenized[index] = repl
+
+        return ' '.join(self.tokenized)
+
+    def perturb(self):
+        count = 1
+        orig_sent = self.sentence
+        for x in range(count):
+            perturb_probs = [.30,.30,.30,.10]
+            perturb_fun = npchoice([self._insert, self._mod_verb, self._replace, self._delete],p=perturb_probs)
+            self.sentence = perturb_fun()
+            self.tokenize()
+        res_sentence = self.sentence
+        self.sentence = self.original_sentence
+        self.tokenize()
+        return res_sentence

critic/critic.py

@@ -0,0 +1,157 @@
+import sys
+import torch
+import random
+import hashlib
+import numpy as np
+from tqdm import tqdm
+from transformers import GPT2Tokenizer, GPT2Model, GPT2LMHeadModel
+import nltk
+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
+
+model_name = 'gpt2'
+tokenizer = GPT2Tokenizer.from_pretrained(model_name)
+tokenizer.pad_token = tokenizer.eos_token
+model = GPT2LMHeadModel.from_pretrained(model_name)
+model.eval()
+#model.cuda()
+model.cpu()
+print (f'Loaded {model_name}')
+
+
+def get_gpt2_loss(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, 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(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, 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, 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, return_string
+
+
+def main():
+    import streamlit as st
+    st.subheader('Exploring Unsupervised Grammatical Error Correction with Transformer-Based Models')
+    sent = st.text_input('Enter a sentence:', value="")
+    if sent != '':
+        st.markdown(f"**Sentence**: {sent}")
+        _,_,_,return_string = gpt2_critic(sent)
+        st.markdown("**Results:**")
+        st.write('\n'.join(return_string))
+
+if __name__ == '__main__':
+    main()

critic/edit_dist_utils.py

@@ -0,0 +1,139 @@
+"""
+Edit distance utils...
+
+Originally by https://worksheets.codalab.org/worksheets/0x8fc01c7fc2b742fdb29c05669f0ad7d2
+"""
+from collections import defaultdict
+import numpy as np
+import random
+import string
+from itertools import permutations
+
+def process_filetype(filetype):
+    insert = (filetype // 1000) % 2 == 1
+    delete = (filetype // 100) % 2 == 1
+    substitute = (filetype // 10) % 2 == 1
+    swap = filetype % 2 == 1
+    return insert, delete, substitute, swap
+
+def get_all_edit_dist_one(word, filetype = 1111, sub_restrict = None):
+    """
+    Allowable edit_dist_one perturbations:
+        1. Insert any lowercase characer at any position other than the start
+        2. Delete any character other than the first one
+        3. Substitute any lowercase character for any other lowercase letter other than the start
+        4. Swap adjacent characters
+    We also include the original word. Filetype determines which of the allowable perturbations to use.
+    """
+    insert, delete, substitute, swap = process_filetype(filetype)
+    #last_mod_pos is last thing you could insert before
+    last_mod_pos = len(word) #- 1
+    ed1 = set()
+    if len(word) <= 2 or word[:1].isupper() or word[:1].isnumeric():
+        return ed1
+    for pos in range(1, last_mod_pos + 1): #can add letters at the end
+        if delete and pos < last_mod_pos:
+            deletion = word[:pos] + word[pos + 1:]
+            ed1.add(deletion)
+        if swap and pos < last_mod_pos - 1:
+            #swapping thing at pos with thing at pos + 1
+            swaped = word[:pos] + word[pos + 1] + word[pos] + word[pos + 2:]
+            ed1.add(swaped)
+        for letter in string.ascii_lowercase: #+"'-": #no need to add '-, as we want to corrupt good to bad
+            if insert:
+                #Insert right after pos - 1
+                insertion = word[:pos] + letter + word[pos:]
+                ed1.add(insertion)
+            can_substitute = sub_restrict is None or letter in sub_restrict[word[pos]]
+            if substitute and pos < last_mod_pos and can_substitute:
+                substitution = word[:pos] + letter + word[pos + 1:]
+                ed1.add(substitution)
+    #Include original word
+    # ed1.add(word)
+    return ed1
+
+def get_all_internal_permutations(word):
+    if len(word) > 10:
+        return set([word])
+    first_char = word[0]
+    last_char = word[-1]
+    internal_chars = word[1:-1]
+    internal_permutations = set()
+    for int_perm in permutations(internal_chars):
+        int_perm_str = ''.join(int_perm)
+        perm = '{}{}{}'.format(first_char, int_perm_str, last_char)
+        internal_permutations.add(perm)
+    return internal_permutations
+
+def sample_random_internal_permutations(word, n_perts = 5):
+    #We try swapping everything with the second character...
+    if len(word) < 4:
+        return set([word])
+    #iterate through positions between second and last
+    perturbations = set()
+    start = word[0]
+    end = word[-1]
+    middle = word[1:-1]
+    for _ in range(n_perts):
+        middle_list = list(middle)
+        random.shuffle(middle_list)
+        mixed_up_middle = ''.join(middle_list)
+        perturbations.add('{}{}{}'.format(start, mixed_up_middle, end))
+    return perturbations
+
+def get_sorted_word(word):
+    if len(word) < 3:
+        sorted_word = word
+    else:
+        sorted_word = '{}{}{}'.format(word[0], ''.join(sorted(word[1:-1])), word[-1])
+    return sorted_word
+
+def get_sorted_word_set(word):
+    if len(word) < 3:
+        sorted_word = word
+    else:
+        sorted_word = '{}{}{}'.format(word[0], ''.join(sorted(word[1:-1])), word[-1])
+    return set([sorted_word])
+
+
+#Used to create agglomerative clusters.
+def preprocess_ed1_neighbors(vocab, sub_restrict = None, filetype = 1111):
+    vocab = set([word.lower() for word in vocab])
+    typo2words = defaultdict(set)
+    for word in vocab:
+        ed1_typos = get_all_edit_dist_one(word, filetype = filetype, sub_restrict = sub_restrict)
+        for typo in ed1_typos:
+            typo2words[typo].add(word)
+
+    word2neighbors = defaultdict(set)
+    for typo in typo2words:
+        for word in typo2words[typo]:
+            word2neighbors[word] = word2neighbors[word].union(typo2words[typo])
+    return word2neighbors
+
+#Used to create agglomerative clusters.
+def ed1_neighbors_mat(vocab, sub_restrict = None, filetype = 1111):
+    vocab = [word.lower() for word in vocab]
+    word2idx = dict([(word, i) for i, word in enumerate(vocab)])
+    word2neighbors = preprocess_ed1_neighbors(vocab, sub_restrict = sub_restrict, filetype = filetype)
+    edges = set()
+    for word in word2neighbors:
+        for neighbor in word2neighbors[word]:
+            edge = [word, neighbor]
+            edge.sort()
+            edge = tuple(edge)
+            edges.add(edge)
+    edge_mat = np.zeros((len(vocab), len(vocab)), dtype = int)
+    for edge in edges:
+        vtx1, vtx2 = edge
+        idx1, idx2 = word2idx[vtx1], word2idx[vtx2]
+        edge_mat[idx1][idx2] = 1
+        edge_mat[idx2][idx1] = 1
+    return edge_mat
+
+
+
+if __name__ == '__main__':
+    while True:
+        word = input("Enter a word: ")
+        print("Total number of possible perturbations: {}".format(len(get_all_edit_dist_one(word))))

critic/perturbations.py

@@ -0,0 +1,144 @@
+"""
+Originally by https://worksheets.codalab.org/worksheets/0x8fc01c7fc2b742fdb29c05669f0ad7d2
+"""
+import json
+import os, sys
+import re
+import random
+import numpy as np
+from random import sample
+from tqdm import tqdm
+from collections import Counter
+
+from critic.edit_dist_utils import get_all_edit_dist_one, sample_random_internal_permutations
+
+
+try:
+    dir_path = os.path.dirname(os.path.realpath(__file__))
+except:
+    dir_path = '.'
+common_typo = json.load(open(f"{dir_path}/common_typo.json"))
+
+random.seed(1234)
+np.random.seed(1234)
+
+
+class RandomPerturbationAttack(object):
+    def __init__(self, attack_type = 'ed1'):
+        self.cache = {} #{word: {0: set(), 1: set(),.. }, ..} #0=swap, 1=substitute, 2=delete, 3=insert
+        self.n_types = 5
+        self.attack_type = attack_type
+    #
+    def sample_perturbations(self, word, n_samples, types):
+        if types is None:
+            type_list = list(range(4)) * (n_samples//4) + list(np.random.choice(self.n_types, n_samples % self.n_types, replace=False))
+        else:
+            type_list = [sample(types,1)[0] for _ in range(n_samples)]
+        type_count = Counter(type_list)
+        perturbations = set()
+        for type in type_count:
+            if type not in self.cache[word]:
+                continue
+            if len(self.cache[word][type]) >= type_count[type]:
+                perturbations.update(set(sample(self.cache[word][type], type_count[type])))
+            else:
+                perturbations.update(self.cache[word][type])
+        return perturbations
+    #
+    def get_perturbations(self, word, n_samples, types=None):
+        if word not in self.cache:
+            self.cache[word] = {}
+            if word[0].islower():
+                for type in range(4):
+                    self.cache[word][type] = get_all_edit_dist_one(word, 10**type)
+                if word in common_typo:
+                    self.cache[word][4] = set(common_typo[word])
+            elif word[0].isupper():
+                if word in common_typo:
+                    self.cache[word][4] = set(common_typo[word])
+        if self.attack_type == 'ed1':
+            perturbations = self.sample_perturbations(word, n_samples, types)
+        else:
+            raise NotImplementedError("Attack type: {} not implemented yet".format(self.attack_type))
+        return perturbations
+    #
+    def name(self):
+        return 'RandomPerturbationAttack'
+
+
+word_attack = RandomPerturbationAttack()
+
+
+def _tokenize(sent):
+    toks = []
+    word_idxs = []
+    for idx, match in enumerate(re.finditer(r'([a-zA-Z]+)|([0-9]+)|.', sent)):
+        tok = match.group(0)
+        toks.append(tok)
+        if len(tok) > 2 and tok.isalpha() and (tok[0].islower()):
+            word_idxs.append(idx)
+    return toks, word_idxs
+
+def _detokenize(toks):
+    return ''.join(toks)
+
+def get_local_neighbors_char_level(sent, max_n_samples=500):
+    words, word_idxs = _tokenize(sent)
+    n_samples = min(len(word_idxs) *20, max_n_samples)
+    sent_perturbations = set()
+    if len(word_idxs) == 0:
+        return sent_perturbations
+    for _ in range(500):
+        word_idx = sample(word_idxs, 1)[0]
+        words_cp = words[:]
+        word_perturbations = list(word_attack.get_perturbations(words_cp[word_idx], n_samples=1))
+        if len(word_perturbations) > 0:
+            words_cp[word_idx] = word_perturbations[0]
+            sent_perturbed = _detokenize(words_cp)
+            if sent_perturbed != sent:
+                sent_perturbations.add(sent_perturbed)
+        if len(sent_perturbations) == n_samples:
+            break
+    #Adding common typos such as 's'
+    for word_idx in word_idxs:
+        words_cp = words[:]
+        word = words_cp[word_idx]
+        if len(word) > 2 and word[0].islower():
+            words_cp[word_idx] = word +'s'
+            sent_perturbed = _detokenize(words_cp)
+            if sent_perturbed != sent:
+                sent_perturbations.add(sent_perturbed)
+            words_cp[word_idx] = word[:-1]
+            sent_perturbed = _detokenize(words_cp)
+            if sent_perturbed != sent:
+                sent_perturbations.add(sent_perturbed)
+    if len(sent_perturbations) > max_n_samples:
+        sent_perturbations = list(sent_perturbations)
+        np.random.shuffle(sent_perturbations)
+        sent_perturbations = set(sent_perturbations[:max_n_samples])
+    return sent_perturbations
+
+
+
+from critic.PIE.word_level_perturb import WordLevelPerturber_all, WordLevelPerturber_refine
+from utils.text_utils import detokenize_sent
+
+def get_local_neighbors_word_level(sent_toked, max_n_samples=500, mode='refine'):
+    """ sent_toked is tokenized by spacy """
+    n_samples = min(len(sent_toked) *20, max_n_samples)
+    orig_sent = ' '.join(sent_toked)
+    orig_sent_detok = detokenize_sent(orig_sent)
+    if mode == 'refine':
+        ptb = WordLevelPerturber_refine(orig_sent)
+    else:
+        ptb = WordLevelPerturber_all(orig_sent)
+    sent_perturbations = set()
+    for _ in range(500):
+        sent_perturbed = ptb.perturb()
+        if sent_perturbed != orig_sent:
+            sent_perturbed_detok = detokenize_sent(sent_perturbed)
+            sent_perturbations.add(sent_perturbed_detok)
+        if len(sent_perturbations) == n_samples:
+            break
+    assert len(sent_perturbations) <= max_n_samples
+    return sent_perturbations, orig_sent_detok

eval_critic/eval_critic.py

@@ -0,0 +1,114 @@
+import re
+import os
+import sys
+import json
+import numpy as np
+import editdistance
+from tqdm import tqdm
+from collections import Counter
+
+sys.path.insert(0, '.')
+from utils.text_utils import detokenize_sent
+from critic.critic import run_gpt2, gpt2_critic
+
+def load_data():
+    data_path = 'eval_critic/eval_data.jsonl'
+    good_sents, bad_sents = [], []
+    for line in open(data_path):
+        obj = json.loads(line)
+        good_sents.append(obj['good'])
+        bad_sents.append(obj['bad'])
+    return good_sents, bad_sents
+
+good_sents, bad_sents = load_data()
+
+
+
+def get_logps(sents):
+    final = []
+    for start in tqdm(range(0, len(sents), 100)):
+        sents_sub = sents[start: start+100]
+        sents_sub_detok = [detokenize_sent(sent) for sent in sents_sub]
+        logps = run_gpt2(sents_sub_detok)
+        assert logps is not None
+        for i in range(len(sents_sub)):
+            final.append({'sent': sents_sub[i], 'sent_detok': sents_sub_detok[i], 'logp': float(logps[i])})
+    return final
+
+def evaluate_logp():
+    """
+    Check whether log p(bad_sent) < log p(good_sent)
+    """
+    good_logps = get_logps(good_sents)
+    bad_logps  = get_logps(bad_sents)
+    accs = []
+    for good, bad in zip(good_logps, bad_logps):
+        accs.append(int(bad['logp'] < good['logp']))
+    avg_acc = float(sum(accs))/len(accs)
+    print (f'log p(bad) < log p(good)? {sum(accs)} / {len(accs)} = {avg_acc:.3f}')
+    return good_logps, bad_logps
+
+good_logps, bad_logps = evaluate_logp()
+# log p(bad) < log p(good)? 555 / 586 = 0.947
+
+
+def compute_metrics(good_accs, bad_accs):
+    goodP = float(sum(good_accs))/(len(bad_accs)-sum(bad_accs)+sum(good_accs))
+    goodR = float(sum(good_accs))/len(good_accs)
+    goodF05 = (1+0.5**2) * float(goodP * goodR)/((0.5**2 * goodP) + goodR)
+    badP = float(sum(bad_accs))/(len(good_accs)-sum(good_accs)+sum(bad_accs))
+    badR = float(sum(bad_accs))/len(bad_accs)
+    badF05 = (1+0.5**2) * float(badP * badR)/((0.5**2 * badP) + badR)
+    print (f'  Good precision = {sum(good_accs)} / {(len(bad_accs)-sum(bad_accs)+sum(good_accs))} = {goodP:.3f}')
+    print (f'  Good recall    = {sum(good_accs)} / {len(good_accs)} = {goodR:.3f}')
+    print (f'  Good F0.5      = {goodF05:.3f}')
+    print (f'  Bad precision  = {sum(bad_accs)} / {(len(good_accs)-sum(good_accs)+sum(bad_accs))} = {badP:.3f}')
+    print (f'  Bad recall     = {sum(bad_accs)} / {len(bad_accs)} = {badR:.3f}')
+    print (f'  Bad F0.5       = {badF05:.3f}')
+    return {'goodP': goodP, 'goodR': goodR, 'goodF05': goodF05, 'badP': badP, 'badR': badR, 'badF05': badF05}
+
+def evaluate_baseline_critic():
+    threshold = np.mean([elm['logp'] for elm in good_logps + bad_logps])
+    good_accs, bad_accs = [], []
+    for obj in good_logps:
+        pred = int(obj['logp'] > threshold)
+        good_accs.append(pred==1)
+    for obj in bad_logps:
+        pred = int(obj['logp'] > threshold)
+        bad_accs.append(pred==0)
+    print ('\nBaseline critic:')
+    stats = compute_metrics(good_accs, bad_accs)
+    json.dump(stats, open('baseline_critic.stats.json', 'w'), indent=2)
+
+evaluate_baseline_critic()
+# Baseline critic:
+#   Good precision = 365 / 668 = 0.546
+#   Good recall    = 365 / 586 = 0.623
+#   Good F0.5      = 0.560
+#   Bad precision  = 283 / 504 = 0.562
+#   Bad recall     = 283 / 586 = 0.483
+#   Bad F0.5       = 0.544
+
+
+def evaluate_LM_Critic():
+    good_accs, bad_accs = [], []
+    for obj in tqdm(good_logps):
+        res = gpt2_critic(obj['sent_detok'], verbose=0, seed=1, n_samples=100, word_level_mode='refine')
+        pred = int(res[0])
+        good_accs.append(pred==1)
+    for obj in tqdm(bad_logps):
+        res = gpt2_critic(obj['sent_detok'], verbose=0, seed=1, n_samples=100, word_level_mode='refine')
+        pred = int(res[0])
+        bad_accs.append(pred==0)
+    print ('\nLM-Critic:')
+    stats = compute_metrics(good_accs, bad_accs)
+    json.dump(stats, open('lm_critic.stats.json', 'w'), indent=2)
+
+evaluate_LM_Critic()
+# LM-Critic: (there is variance due to the randomness of sampling, some variation in GPT2 return score)
+#   Good precision = 446 / 654 = 0.682
+#   Good recall    = 446 / 586 = 0.761
+#   Good F0.5      = 0.696
+#   Bad precision  = 378 / 518 = 0.730
+#   Bad recall     = 378 / 586 = 0.645
+#   Bad F0.5       = 0.711

gec/download_data.sh

@@ -0,0 +1,43 @@
+conda activate errant200
+
+
+######################## Set up benckmarks ########################
+mkdir -p benchmarks
+cd benchmarks
+
+#Prepare CoNLL2014
+wget https://www.comp.nus.edu.sg/~nlp/conll14st/conll14st-test-data.tar.gz
+tar -xf conll14st-test-data.tar.gz
+python3 scripts/get_orig_from_m2.py conll14st-test-data/noalt/official-2014.combined.m2 \
+      -out conll14st-test-data/noalt/official-2014.combined.orig.txt
+
+
+#Prepare BEA2019
+wget https://www.cl.cam.ac.uk/research/nl/bea2019st/data/wi+locness_v2.1.bea19.tar.gz
+tar -xf wi+locness_v2.1.bea19.tar.gz
+mv wi+locness wi+locness_v2.1.bea19
+python3 scripts/get_orig_from_m2.py wi+locness_v2.1.bea19/m2/ABCN.dev.gold.bea19.m2 \
+      -out wi+locness_v2.1.bea19/m2/ABCN.dev.bea19.orig.txt
+
+
+#Prepare GMEG-wiki and -yahoo
+git clone https://github.com/grammarly/GMEG.git
+root=GMEG/data/test/wiki
+errant_parallel -orig $root/source \
+                -cor  $root/ref0 $root/ref1 $root/ref2 $root/ref3 \
+                -out  $root/ref.m2
+
+root=GMEG/data/test/yahoo
+errant_parallel -orig $root/source \
+                -cor  $root/ref0 $root/ref1 $root/ref2 $root/ref3 \
+                -out  $root/ref.m2
+
+
+#Download M2 scorer
+git clone https://github.com/nusnlp/m2scorer.git
+
+
+######################## Download training data ########################
+cd ../
+wget https://nlp.stanford.edu/projects/myasu/LM-Critic/data.zip
+unzip data.zip

gec/scripts/get_corr_from_m2.py

@@ -0,0 +1,35 @@
+import argparse
+
+# Apply the edits of a single annotator to generate the corrected sentences.
+def main(args):
+	m2 = open(args.m2_file).read().strip().split("\n\n")
+	out = open(args.out, "w")
+	# Do not apply edits with these error types
+	skip = {"noop", "UNK", "Um"}
+
+	for sent in m2:
+		sent = sent.split("\n")
+		cor_sent = sent[0].split()[1:] # Ignore "S "
+		edits = sent[1:]
+		offset = 0
+		for edit in edits:
+			edit = edit.split("|||")
+			if edit[1] in skip: continue # Ignore certain edits
+			coder = int(edit[-1])
+			if coder != args.id: continue # Ignore other coders
+			span = edit[0].split()[1:] # Ignore "A "
+			start = int(span[0])
+			end = int(span[1])
+			cor = edit[2].split()
+			cor_sent[start+offset:end+offset] = cor
+			offset = offset-(end-start)+len(cor)
+		out.write(" ".join(cor_sent)+"\n")
+
+if __name__ == "__main__":
+	# Define and parse program input
+	parser = argparse.ArgumentParser()
+	parser.add_argument("m2_file", help="The path to an input m2 file.")
+	parser.add_argument("-out", help="A path to where we save the output corrected text file.", required=True)
+	parser.add_argument("-id", help="The id of the target annotator in the m2 file.", type=int, default=0)
+	args = parser.parse_args()
+	main(args)

gec/scripts/get_orig_from_m2.py

@@ -0,0 +1,22 @@
+import argparse
+
+# Apply the edits of a single annotator to generate the corrected sentences.
+def main(args):
+	m2 = open(args.m2_file).read().strip().split("\n\n")
+	out = open(args.out, "w")
+	# Do not apply edits with these error types
+	skip = {"noop", "UNK", "Um"}
+
+	for sent in m2:
+		sent = sent.split("\n")
+		orig_sent = sent[0].split()[1:] # Ignore "S "
+		out.write(" ".join(orig_sent)+"\n")
+
+if __name__ == "__main__":
+	# Define and parse program input
+	parser = argparse.ArgumentParser()
+	parser.add_argument("m2_file", help="The path to an input m2 file.")
+	parser.add_argument("-out", help="A path to where we save the output corrected text file.", required=True)
+	parser.add_argument("-id", help="The id of the target annotator in the m2 file.", type=int, default=0)
+	args = parser.parse_args()
+	main(args)

gec/scripts/parse_errant_output.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+
+import sys
+import json
+
+for i, line in enumerate(sys.stdin):
+    if i == 3:
+        nums = line.split()
+        P, R, F = nums[3:6]
+        json.dump({'precision': float(P), 'recall': float(R), 'F0.5': float(F)}, open('stats.json', 'w'), indent=2)
+        break

gec/scripts/parse_m2_output.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+
+import sys
+import json
+
+scores = []
+for i, line in enumerate(sys.stdin):
+    score = line.split(':')[1].strip()
+    scores.append(float(score))
+
+json.dump({'precision': scores[0], 'recall': scores[1], 'F0.5': scores[2]}, open('stats.json', 'w'), indent=2)

gec/src/run-round0.sh

@@ -0,0 +1,77 @@
+exit 0;
+################################################################################
+# run the following commands one by one in the `gec/` directory of the repo
+################################################################################
+export CUDA_VISIBLE_DEVICES=0
+conda activate lm-critic
+
+############### Train the fixer ###############
+dt=`date '+%Y%m%d_%H%M%S'`
+outdir=data/round0__synthetic/model-fixer__${dt}
+mkdir -p $outdir
+python3.8 -u src/run_seq2seq.py \
+    --model_name_or_path facebook/bart-base --task summarization --text_column bad_detoked --summary_column good_detoked \
+    --do_train --num_train_epochs 1 --train_file data/round0__synthetic/synthetic_paired_data_9M.json \
+    --preprocessing_num_workers 20 --overwrite_output_dir --output_dir $outdir --predict_with_generate --fp16 \
+    --per_device_train_batch_size 64 --gradient_accumulation_steps 8 --max_source_length 64 --max_target_length 64 \
+    --logging_first_step --logging_steps 20 --save_steps 2000 \
+  |& tee $outdir/log.txt
+
+
+
+############### Run the fixer on benchmarks ###############
+model_path=data/round0__synthetic/model-fixer
+
+#BEA2019
+python src/run_fixer.py -m $model_path -i benchmarks/wi+locness_v2.1.bea19/m2/ABCN.dev.bea19.orig.txt -o $model_path/predictions/bea19dev.out.txt --bea19
+#CoNLL2014
+python src/run_fixer.py -m $model_path -i benchmarks/conll14st-test-data/noalt/official-2014.combined.orig.txt -o $model_path/predictions/conll14.out.txt
+#GMEG-wiki
+python src/run_fixer.py -m $model_path -i benchmarks/GMEG/data/test/wiki/source -o $model_path/predictions/gmeg.wiki.out.txt
+#GMEG-yahoo
+python src/run_fixer.py -m $model_path -i benchmarks/GMEG/data/test/yahoo/source -o $model_path/predictions/gmeg.yahoo.out.txt
+
+
+
+############### Evaluate the fixer outputs ###############
+#CoNLL2014
+python2 benchmarks/m2scorer/scripts/m2scorer.py $model_path/predictions/conll14.out.txt \
+    benchmarks/conll14st-test-data/noalt/official-2014.combined.m2 | tee $model_path/predictions/conll14.eval.txt
+# Precision   : 0.5922
+# Recall      : 0.2920
+# F_0.5       : 0.4912
+
+
+#BEA2019 and GMEG uses errant scorer, which needs its own environment
+conda deactivate
+conda activate errant200
+
+#BEA2019
+errant_parallel -orig benchmarks/wi+locness_v2.1.bea19/m2/ABCN.dev.bea19.orig.txt \
+                -cor $model_path/predictions/bea19dev.out.txt \
+                -out $model_path/predictions/bea19dev.outm2.txt && \
+errant_compare  -hyp $model_path/predictions/bea19dev.outm2.txt -ref benchmarks/wi+locness_v2.1.bea19/m2/ABCN.dev.gold.bea19.m2 | tee $model_path/predictions/bea19dev.eval.txt
+# =========== Span-Based Correction ============
+# TP	FP	FN	Prec	Rec	F0.5
+# 1337	1686	6124	0.4423	0.1792	0.3419
+# ==============================================
+
+#GEMG-wiki
+errant_parallel -orig benchmarks/GMEG/data/test/wiki/source \
+                -cor $model_path/predictions/gmeg.wiki.out.txt \
+                -out $model_path/predictions/gmeg.wiki.outm2.txt && \
+errant_compare  -hyp $model_path/predictions/gmeg.wiki.outm2.txt -ref benchmarks/GMEG/data/test/wiki/ref.m2 | tee $model_path/predictions/gmeg.wiki.eval.txt
+# =========== Span-Based Correction ============
+# TP	FP	FN	Prec	Rec	F0.5
+# 352	323	973	0.5215	0.2657	0.4373
+# ==============================================
+
+#GEMG-yahoo
+errant_parallel -orig benchmarks/GMEG/data/test/yahoo/source \
+                -cor $model_path/predictions/gmeg.yahoo.out.txt \
+                -out $model_path/predictions/gmeg.yahoo.outm2.txt && \
+errant_compare  -hyp $model_path/predictions/gmeg.yahoo.outm2.txt -ref benchmarks/GMEG/data/test/yahoo/ref.m2 | tee $model_path/predictions/gmeg.yahoo.eval.txt
+# =========== Span-Based Correction ============
+# TP	FP	FN	Prec	Rec	F0.5
+# 241	301	411	0.4446	0.3696	0.4273
+# ==============================================

gec/src/run-round1.sh

@@ -0,0 +1,75 @@
+exit 0;
+################################################################################
+# run the following commands one by one in the `gec/` directory of the repo
+################################################################################
+export CUDA_VISIBLE_DEVICES=0
+conda activate lm-critic
+
+############### Train the fixer ###############
+dt=`date '+%Y%m%d_%H%M%S'`
+outdir=data/round1__BIFI/model-fixer__${dt}
+mkdir -p $outdir
+python3.8 -u src/run_seq2seq.py \
+    --model_name_or_path facebook/bart-base --task summarization --text_column bad_detoked --summary_column good_detoked \
+    --do_train --num_train_epochs 1 --train_file data/round1__BIFI/BIFI_paired_data_9M.json \
+    --preprocessing_num_workers 20 --overwrite_output_dir --output_dir $outdir --predict_with_generate --fp16 \
+    --per_device_train_batch_size 64 --gradient_accumulation_steps 8 --max_source_length 64 --max_target_length 64 \
+    --logging_first_step --logging_steps 20 --save_steps 2000 \
+  |& tee $outdir/log.txt
+
+
+
+############### Run the fixer on benchmarks ###############
+model_path=data/round1__BIFI/model-fixer
+#BEA2019
+python src/run_fixer.py -m $model_path -i benchmarks/wi+locness_v2.1.bea19/m2/ABCN.dev.bea19.orig.txt -o $model_path/predictions/bea19dev.out.txt --bea19
+#CoNLL2014
+python src/run_fixer.py -m $model_path -i benchmarks/conll14st-test-data/noalt/official-2014.combined.orig.txt -o $model_path/predictions/conll14.out.txt
+#GMEG-wiki
+python src/run_fixer.py -m $model_path -i benchmarks/GMEG/data/test/wiki/source -o $model_path/predictions/gmeg.wiki.out.txt
+#GMEG-yahoo
+python src/run_fixer.py -m $model_path -i benchmarks/GMEG/data/test/yahoo/source -o $model_path/predictions/gmeg.yahoo.out.txt
+
+
+
+############### Evaluate the fixer outputs ###############
+#CoNLL2014
+python2 benchmarks/m2scorer/scripts/m2scorer.py $model_path/predictions/conll14.out.txt \
+    benchmarks/conll14st-test-data/noalt/official-2014.combined.m2 | tee $model_path/predictions/conll14.eval.txt
+# Precision   : 0.6444
+# Recall      : 0.3569
+# F_0.5       : 0.5550
+
+#BEA2019 and GMEG uses errant scorer, which needs its own environment
+conda deactivate
+conda activate errant200
+
+#BEA2019
+errant_parallel -orig benchmarks/wi+locness_v2.1.bea19/m2/ABCN.dev.bea19.orig.txt \
+                -cor $model_path/predictions/bea19dev.out.txt \
+                -out $model_path/predictions/bea19dev.outm2.txt && \
+errant_compare  -hyp $model_path/predictions/bea19dev.outm2.txt -ref benchmarks/wi+locness_v2.1.bea19/m2/ABCN.dev.gold.bea19.m2 | tee $model_path/predictions/bea19dev.eval.txt
+# =========== Span-Based Correction ============
+# TP	FP	FN	Prec	Rec	F0.5
+# 1848	1733	5613	0.5161	0.2477	0.4241
+# ==============================================
+
+#GEMG-wiki
+errant_parallel -orig benchmarks/GMEG/data/test/wiki/source \
+                -cor $model_path/predictions/gmeg.wiki.out.txt \
+                -out $model_path/predictions/gmeg.wiki.outm2.txt && \
+errant_compare  -hyp $model_path/predictions/gmeg.wiki.outm2.txt -ref benchmarks/GMEG/data/test/wiki/ref.m2 | tee $model_path/predictions/gmeg.wiki.eval.txt
+# =========== Span-Based Correction ============
+# TP	FP	FN	Prec	Rec	F0.5
+# 468	339	925	0.5799	0.336	0.5064
+# ==============================================
+
+#GEMG-yahoo
+errant_parallel -orig benchmarks/GMEG/data/test/yahoo/source \
+                -cor $model_path/predictions/gmeg.yahoo.out.txt \
+                -out $model_path/predictions/gmeg.yahoo.outm2.txt && \
+errant_compare  -hyp $model_path/predictions/gmeg.yahoo.outm2.txt -ref benchmarks/GMEG/data/test/yahoo/ref.m2 | tee $model_path/predictions/gmeg.yahoo.eval.txt
+# =========== Span-Based Correction ============
+# TP	FP	FN	Prec	Rec	F0.5
+# 382	329	428	0.5373	0.4716	0.5227
+# ==============================================

gec/src/run_fixer.py

@@ -0,0 +1,87 @@
+import os
+import sys
+import json
+import torch
+import argparse
+from tqdm import tqdm
+from transformers import BartForConditionalGeneration, BartTokenizer
+
+sys.path.insert(0, '..')
+from utils.text_utils import detokenize_sent
+from utils.spacy_tokenizer import spacy_tokenize_gec, spacy_tokenize_bea19
+
+parser = argparse.ArgumentParser()
+parser.add_argument('-m', '--model_path')
+parser.add_argument('-i', '--input_path')
+parser.add_argument('-o', '--output_path')
+parser.add_argument('--bea19', action='store_true')
+args = parser.parse_args()
+
+
+tokenizer = BartTokenizer.from_pretrained("facebook/bart-base")
+model = BartForConditionalGeneration.from_pretrained(args.model_path, force_bos_token_to_be_generated=True)
+model.eval()
+model.cuda()
+
+
+def run_model(sents):
+    num_ret_seqs = 10
+    inp_max_len = 66
+    batch = [tokenizer(s, return_tensors='pt', padding='max_length', max_length=inp_max_len) for s in sents]
+    oidx2bidx = {} #original index to final batch index
+    final_batch = []
+    for oidx, elm in enumerate(batch):
+        if elm['input_ids'].size(1) <= inp_max_len:
+            oidx2bidx[oidx] = len(final_batch)
+            final_batch.append(elm)
+    batch = {key: torch.cat([elm[key] for elm in final_batch], dim=0) for key in final_batch[0]}
+    with torch.no_grad():
+        generated_ids = model.generate(batch['input_ids'].cuda(),
+                                attention_mask=batch['attention_mask'].cuda(),
+                                num_beams=10, num_return_sequences=num_ret_seqs, max_length=65)
+    _out = tokenizer.batch_decode(generated_ids.detach().cpu(), skip_special_tokens=True)
+    outs = []
+    for i in range(0, len(_out), num_ret_seqs):
+        outs.append(_out[i:i+num_ret_seqs])
+    final_outs = [[sents[oidx]] if oidx not in oidx2bidx else outs[oidx2bidx[oidx]] for oidx in range(len(sents))]
+    return final_outs
+
+
+def run_for_wiki_yahoo_conll():
+    sents = [detokenize_sent(l.strip()) for l in open(args.input_path)]
+    b_size = 40
+    outs = []
+    for j in tqdm(range(0, len(sents), b_size)):
+        sents_batch = sents[j:j+b_size]
+        outs_batch = run_model(sents_batch)
+        for sent, preds in zip(sents_batch, outs_batch):
+            preds_detoked = [detokenize_sent(pred) for pred in preds]
+            preds = [' '.join(spacy_tokenize_gec(pred)) for pred in preds_detoked]
+            outs.append({'src': sent, 'preds': preds})
+    os.system('mkdir -p {}'.format(os.path.dirname(args.output_path)))
+    with open(args.output_path, 'w') as outf:
+        for out in outs:
+            print (out['preds'][0], file=outf)
+
+
+def run_for_bea19():
+    sents = [detokenize_sent(l.strip()) for l in open(args.input_path)]
+    b_size = 40
+    outs = []
+    for j in tqdm(range(0, len(sents), b_size)):
+        sents_batch = sents[j:j+b_size]
+        outs_batch = run_model(sents_batch)
+        for sent, preds in zip(sents_batch, outs_batch):
+            preds_detoked = [detokenize_sent(pred) for pred in preds]
+            preds = [' '.join(spacy_tokenize_bea19(pred)) for pred in preds_detoked]
+            outs.append({'src': sent, 'preds': preds})
+    os.system('mkdir -p {}'.format(os.path.dirname(args.output_path)))
+    with open(args.output_path, 'w') as outf:
+        for out in outs:
+            print (out['preds'][0], file=outf)
+
+
+if args.bea19:
+    run_for_bea19()
+else:
+    run_for_wiki_yahoo_conll()

gec/src/run_seq2seq.py

@@ -0,0 +1,537 @@
+# coding=utf-8
+# Copyright The HuggingFace Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for sequence to sequence.
+"""
+# You can also adapt this script on your own sequence to sequence task. Pointers for this are left as comments.
+
+import logging
+import os
+import re
+import sys
+from dataclasses import dataclass, field
+from typing import Optional
+
+import numpy as np
+from datasets import load_dataset, load_metric
+
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    DataCollatorForSeq2Seq,
+    HfArgumentParser,
+    MBartTokenizer,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+    default_data_collator,
+    set_seed,
+)
+from transformers.trainer_utils import get_last_checkpoint, is_main_process
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    model_revision: str = field(
+        default="main",
+        metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."},
+    )
+    use_auth_token: bool = field(
+        default=False,
+        metadata={
+            "help": "Will use the token generated when running `transformers-cli login` (necessary to use this script "
+            "with private models)."
+        },
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    task: str = field(
+        default="summarization",
+        metadata={
+            "help": "The name of the task, should be summarization (or summarization_{dataset} for evaluating "
+            "pegasus) or translation (or translation_{xx}_to_{yy})."
+        },
+    )
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    text_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "The name of the column in the datasets containing the full texts (for summarization)."},
+    )
+    summary_column: Optional[str] = field(
+        default=None,
+        metadata={"help": "The name of the column in the datasets containing the summaries (for summarization)."},
+    )
+    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_source_length: Optional[int] = field(
+        default=1024,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    max_target_length: Optional[int] = field(
+        default=128,
+        metadata={
+            "help": "The maximum total sequence length for target text after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    val_max_target_length: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded. Will default to `max_target_length`."
+            "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
+            "during ``evaluate`` and ``predict``."
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad all samples to model maximum sentence length. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
+            "efficient on GPU but very bad for TPU."
+        },
+    )
+    max_train_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of training examples to this "
+            "value if set."
+        },
+    )
+    max_val_samples: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "For debugging purposes or quicker training, truncate the number of validation examples to this "
+            "value if set."
+        },
+    )
+    source_lang: Optional[str] = field(default=None, metadata={"help": "Source language id for translation."})
+    target_lang: Optional[str] = field(default=None, metadata={"help": "Target language id for translation."})
+    eval_beams: Optional[int] = field(default=None, metadata={"help": "Number of beams to use for evaluation."})
+    ignore_pad_token_for_loss: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to ignore the tokens corresponding to padded labels in the loss computation or not."
+        },
+    )
+    source_prefix: Optional[str] = field(
+        default=None, metadata={"help": "A prefix to add before every source text (useful for T5 models)."}
+    )
+
+    def __post_init__(self):
+        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
+            raise ValueError("Need either a dataset name or a training/validation file.")
+        else:
+            if self.train_file is not None:
+                extension = self.train_file.split(".")[-1]
+                assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
+            if self.validation_file is not None:
+                extension = self.validation_file.split(".")[-1]
+                assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
+        if not self.task.startswith("summarization") and not self.task.startswith("translation"):
+            raise ValueError(
+                "`task` should be summarization, summarization_{dataset}, translation or translation_{xx}_to_{yy}."
+            )
+        if self.val_max_target_length is None:
+            self.val_max_target_length = self.max_target_length
+
+
+summarization_name_mapping = {
+    "amazon_reviews_multi": ("review_body", "review_title"),
+    "big_patent": ("description", "abstract"),
+    "cnn_dailymail": ("article", "highlights"),
+    "orange_sum": ("text", "summary"),
+    "pn_summary": ("article", "summary"),
+    "psc": ("extract_text", "summary_text"),
+    "samsum": ("dialogue", "summary"),
+    "thaisum": ("body", "summary"),
+    "xglue": ("news_body", "news_title"),
+    "xsum": ("document", "summary"),
+    "wiki_summary": ("article", "highlights"),
+}
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    # Detecting last checkpoint.
+    last_checkpoint = None
+    if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir:
+        last_checkpoint = get_last_checkpoint(training_args.output_dir)
+        if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0:
+            raise ValueError(
+                f"Output directory ({training_args.output_dir}) already exists and is not empty. "
+                "Use --overwrite_output_dir to overcome."
+            )
+        elif last_checkpoint is not None:
+            logger.info(
+                f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change "
+                "the `--output_dir` or add `--overwrite_output_dir` to train from scratch."
+            )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        handlers=[logging.StreamHandler(sys.stdout)],
+    )
+    logger.setLevel(logging.INFO if is_main_process(training_args.local_rank) else logging.WARN)
+
+    # Log on each process the small summary:
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+    # Set the verbosity to info of the Transformers logger (on main process only):
+    if is_main_process(training_args.local_rank):
+        transformers.utils.logging.set_verbosity_info()
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    # Get the datasets: you can either provide your own CSV/JSON training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files in the summarization task, this script will use the first column for the full texts and the
+    # second column for the summaries (unless you specify column names for this with the `text_column` and
+    # `summary_column` arguments).
+    # For translation, only JSON files are supported, with one field named "translation" containing two keys for the
+    # source and target languages (unless you adapt what follows).
+    #
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if data_args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
+    else:
+        data_files = {}
+        if data_args.train_file is not None:
+            data_files["train"] = data_args.train_file
+            extension = data_args.train_file.split(".")[-1]
+        if data_args.validation_file is not None:
+            data_files["validation"] = data_args.validation_file
+            extension = data_args.validation_file.split(".")[-1]
+        datasets = load_dataset(extension, data_files=data_files)
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    config = AutoConfig.from_pretrained(
+        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_args.model_name_or_path,
+        from_tf=bool(".ckpt" in model_args.model_name_or_path),
+        config=config,
+        cache_dir=model_args.cache_dir,
+        revision=model_args.model_revision,
+        use_auth_token=True if model_args.use_auth_token else None,
+    )
+
+    # Set decoder_start_token_id
+    if model.config.decoder_start_token_id is None and isinstance(tokenizer, MBartTokenizer):
+        model.config.decoder_start_token_id = tokenizer.lang_code_to_id[data_args.target_lang]
+    if model.config.decoder_start_token_id is None:
+        raise ValueError("Make sure that `config.decoder_start_token_id` is correctly defined")
+
+    # Get the default prefix if None is passed.
+    if data_args.source_prefix is None:
+        task_specific_params = model.config.task_specific_params
+        if task_specific_params is not None:
+            prefix = task_specific_params.get("prefix", "")
+        else:
+            prefix = ""
+    else:
+        prefix = data_args.source_prefix
+
+    # Preprocessing the datasets.
+    # We need to tokenize inputs and targets.
+    if training_args.do_train:
+        column_names = datasets["train"].column_names
+    else:
+        column_names = datasets["validation"].column_names
+
+    # For translation we set the codes of our source and target languages (only useful for mBART, the others will
+    # ignore those attributes).
+    if data_args.task.startswith("translation"):
+        if data_args.source_lang is not None:
+            tokenizer.src_lang = data_args.source_lang
+        if data_args.target_lang is not None:
+            tokenizer.tgt_lang = data_args.target_lang
+
+    # To serialize preprocess_function below, each of those four variables needs to be defined (even if we won't use
+    # them all).
+    source_lang, target_lang, text_column, summary_column = None, None, None, None
+
+    if data_args.task.startswith("summarization"):
+        # Get the column names for input/target.
+        dataset_columns = summarization_name_mapping.get(data_args.dataset_name, None)
+        if data_args.text_column is None:
+            text_column = dataset_columns[0] if dataset_columns is not None else column_names[0]
+        else:
+            text_column = data_args.text_column
+        if data_args.summary_column is None:
+            summary_column = dataset_columns[1] if dataset_columns is not None else column_names[1]
+        else:
+            summary_column = data_args.summary_column
+    else:
+        # Get the language codes for input/target.
+        lang_search = re.match("translation_([a-z]+)_to_([a-z]+)", data_args.task)
+        if data_args.source_lang is not None:
+            source_lang = data_args.source_lang.split("_")[0]
+        else:
+            assert (
+                lang_search is not None
+            ), "Provide a source language via --source_lang or rename your task 'translation_xx_to_yy'."
+            source_lang = lang_search.groups()[0]
+
+        if data_args.target_lang is not None:
+            target_lang = data_args.target_lang.split("_")[0]
+        else:
+            assert (
+                lang_search is not None
+            ), "Provide a target language via --target_lang or rename your task 'translation_xx_to_yy'."
+            target_lang = lang_search.groups()[1]
+
+    # Temporarily set max_target_length for training.
+    max_target_length = data_args.max_target_length
+    padding = "max_length" if data_args.pad_to_max_length else False
+
+    def preprocess_function(examples):
+        if data_args.task.startswith("translation"):
+            inputs = [ex[source_lang] for ex in examples["translation"]]
+            targets = [ex[target_lang] for ex in examples["translation"]]
+        else:
+            inputs = examples[text_column]
+            targets = examples[summary_column]
+        inputs = [prefix + inp for inp in inputs]
+        model_inputs = tokenizer(inputs, max_length=data_args.max_source_length, padding=padding, truncation=True)
+
+        # Setup the tokenizer for targets
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, max_length=max_target_length, padding=padding, truncation=True)
+
+        # If we are padding here, replace all tokenizer.pad_token_id in the labels by -100 when we want to ignore
+        # padding in the loss.
+        if padding == "max_length" and data_args.ignore_pad_token_for_loss:
+            labels["input_ids"] = [
+                [(l if l != tokenizer.pad_token_id else -100) for l in label] for label in labels["input_ids"]
+            ]
+
+        model_inputs["labels"] = labels["input_ids"]
+        return model_inputs
+
+    if training_args.do_train:
+        train_dataset = datasets["train"]
+        if data_args.max_train_samples is not None:
+            train_dataset = train_dataset.select(range(data_args.max_train_samples))
+        train_dataset = train_dataset.map(
+            preprocess_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
+    if training_args.do_eval:
+        max_target_length = data_args.val_max_target_length
+        eval_dataset = datasets["validation"]
+        if data_args.max_val_samples is not None:
+            eval_dataset = eval_dataset.select(range(data_args.max_val_samples))
+        eval_dataset = eval_dataset.map(
+            preprocess_function,
+            batched=True,
+            num_proc=data_args.preprocessing_num_workers,
+            remove_columns=column_names,
+            load_from_cache_file=not data_args.overwrite_cache,
+        )
+
+    # Data collator
+    label_pad_token_id = -100 if data_args.ignore_pad_token_for_loss else tokenizer.pad_token_id
+    if data_args.pad_to_max_length:
+        data_collator = default_data_collator
+    else:
+        data_collator = DataCollatorForSeq2Seq(
+            tokenizer,
+            label_pad_token_id=label_pad_token_id,
+            pad_to_multiple_of=8 if training_args.fp16 else None,
+        )
+
+    # Metric
+    metric_name = "rouge" if data_args.task.startswith("summarization") else "sacrebleu"
+    metric = load_metric(metric_name)
+
+    def compute_metrics(eval_preds):
+        preds, labels = eval_preds
+        if isinstance(preds, tuple):
+            preds = preds[0]
+        decoded_preds = tokenizer.batch_decode(preds, skip_special_tokens=True)
+        if data_args.ignore_pad_token_for_loss:
+            # Replace -100 in the labels as we can't decode them.
+            labels = np.where(labels != -100, labels, tokenizer.pad_token_id)
+        decoded_labels = tokenizer.batch_decode(labels, skip_special_tokens=True)
+
+        # Some simple post-processing
+        decoded_preds = [pred.strip() for pred in decoded_preds]
+        decoded_labels = [label.strip() for label in decoded_labels]
+        if metric_name == "sacrebleu":
+            decoded_labels = [[label] for label in decoded_labels]
+
+        result = metric.compute(predictions=decoded_preds, references=decoded_labels)
+
+        # Extract a few results from ROUGE
+        if metric_name == "rouge":
+            result = {key: value.mid.fmeasure * 100 for key, value in result.items()}
+        else:
+            result = {"bleu": result["score"]}
+
+        prediction_lens = [np.count_nonzero(pred != tokenizer.pad_token_id) for pred in preds]
+        result["gen_len"] = np.mean(prediction_lens)
+
+        return result
+
+    # Initialize our Trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset if training_args.do_train else None,
+        eval_dataset=eval_dataset if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics if training_args.predict_with_generate else None,
+    )
+
+    # Training
+    if training_args.do_train:
+        if last_checkpoint is not None:
+            checkpoint = last_checkpoint
+        elif os.path.isdir(model_args.model_name_or_path):
+            checkpoint = model_args.model_name_or_path
+        else:
+            checkpoint = None
+        train_result = trainer.train(resume_from_checkpoint=checkpoint)
+        trainer.save_model()  # Saves the tokenizer too for easy upload
+
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
+    # Evaluation
+    results = {}
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+
+        results = trainer.evaluate()
+
+        output_eval_file = os.path.join(training_args.output_dir, "eval_results_seq2seq.txt")
+        if trainer.is_world_process_zero():
+            with open(output_eval_file, "w") as writer:
+                logger.info("***** Eval results *****")
+                for key, value in sorted(results.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+    return results
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -0,0 +1,9 @@
+datasets==1.3.0
+editdistance==0.6.0
+nltk==3.7
+numpy==1.22.3
+spacy==3.0.5
+streamlit==1.9.0
+torch==1.11.0
+tqdm==4.49.0
+transformers==4.3.3

utils/spacy_tokenizer.py

@@ -0,0 +1,62 @@
+import spacy
+from spacy.tokenizer import Tokenizer
+from spacy.lang.char_classes import ALPHA, ALPHA_LOWER, ALPHA_UPPER, CONCAT_QUOTES, LIST_ELLIPSES, LIST_ICONS, HYPHENS
+from spacy.util import compile_infix_regex
+from spacy.lang.en import English
+nlp = English()
+
+def get_tokenizer_gec(nlp):
+    infixes = (
+        LIST_ELLIPSES
+        + LIST_ICONS
+        + [
+            r"(?<=[0-9])[+\-\*^](?=[0-9-])",
+            r"(?<=[{al}{q}])\.(?=[{au}{q}])".format(
+                al=ALPHA_LOWER, au=ALPHA_UPPER, q=CONCAT_QUOTES
+            ),
+            r"(?<=[{a}]),(?=[{a}])".format(a=ALPHA),
+            #r"(?<=[{a}])(?:{h})(?=[{a}])".format(a=ALPHA, h=HYPHENS),
+            r"(?<=[{a}0-9])[:<>=/](?=[{a}])".format(a=ALPHA),
+        ]
+    )
+    infix_re = compile_infix_regex(infixes)
+    return Tokenizer(nlp.vocab, prefix_search=nlp.tokenizer.prefix_search,
+                                suffix_search=nlp.tokenizer.suffix_search,
+                                infix_finditer=infix_re.finditer,
+                                token_match=nlp.tokenizer.token_match,
+                                rules=nlp.Defaults.tokenizer_exceptions)
+
+
+def get_tokenizer_bea19(nlp):
+    infixes = (
+        LIST_ELLIPSES
+        + LIST_ICONS
+        + [
+            r"(?<=[0-9])[+\-\*^](?=[0-9-])",
+            r"(?<=[{al}{q}])\.(?=[{au}{q}])".format(
+                al=ALPHA_LOWER, au=ALPHA_UPPER, q=CONCAT_QUOTES
+            ),
+            r"(?<=[{a}]),(?=[{a}])".format(a=ALPHA),
+            r"(?<=[{a}])(?:{h})(?=[{a}])".format(a=ALPHA, h=HYPHENS),
+            r"(?<=[{a}0-9])[:<>=/](?=[{a}])".format(a=ALPHA),
+        ]
+    )
+    infix_re = compile_infix_regex(infixes)
+    return Tokenizer(nlp.vocab, prefix_search=nlp.tokenizer.prefix_search,
+                                suffix_search=nlp.tokenizer.suffix_search,
+                                infix_finditer=infix_re.finditer,
+                                token_match=nlp.tokenizer.token_match,
+                                rules=nlp.Defaults.tokenizer_exceptions)
+
+
+tokenizer_gec = get_tokenizer_gec(nlp)
+tokenizer_bea19 = get_tokenizer_bea19(nlp)
+
+
+def spacy_tokenize_gec(text):
+    nlp.tokenizer = tokenizer_gec
+    return [str(w) for w in nlp(text)]
+
+def spacy_tokenize_bea19(text):
+    nlp.tokenizer = tokenizer_bea19
+    return [str(w) for w in nlp(text)]

utils/text_utils.py

@@ -0,0 +1,65 @@
+import re
+from nltk import sent_tokenize, word_tokenize
+from nltk.tokenize.treebank import TreebankWordDetokenizer
+detokenizer = TreebankWordDetokenizer()
+
+def handle_dounble_quote(sent):
+    cur_str = ''
+    exp_left = True
+    ignore_space = False
+    for char in sent:
+        if char == '"':
+            if exp_left: #this is a left "
+                cur_str = cur_str.rstrip() + ' "'
+                exp_left = (not exp_left)
+                ignore_space = True
+            else: #this is a right "
+                cur_str = cur_str.rstrip() + '" '
+                exp_left = (not exp_left)
+                ignore_space = False
+        else:
+            if ignore_space: #expecting right
+                if char == ' ':
+                    continue
+                else:
+                    cur_str = cur_str + char
+                    ignore_space = False
+            else:
+                cur_str = cur_str + char
+    cur_str = cur_str.strip()
+    cur_str = re.sub(r'[ ]+', ' ', cur_str)
+    return cur_str
+
+def postprocess_space(sent):
+    sent = re.sub(r'[ ]+\.', '.', sent)
+    sent = re.sub(r'[ ]+,', ',', sent)
+    sent = re.sub(r'[ ]+!', '!', sent)
+    sent = re.sub(r'[ ]+\?', '?', sent)
+    sent = re.sub(r'\([ ]+', '(', sent)
+    sent = re.sub(r'[ ]+\)', ')', sent)
+    sent = re.sub(r' \'s( |\.|,|!|\?)', r"'s\1", sent)
+    sent = re.sub(r'n \'t( |\.|,|!|\?)', r"n't\1", sent)
+    return sent
+
+def detokenize_sent(sent):
+    #Clean raw sent
+    sent = re.sub(r'\' s ', '\'s ', sent)
+    toks = sent.split()
+    if len([1 for t in toks if t=="'"]) % 2 == 0:
+        toks = ['"' if t=="'" else t for t in toks]
+    sent = ' '.join(toks)
+    #
+    sents = sent_tokenize(sent)
+    final_sents = []
+    for _sent in sents:
+        _sent = detokenizer.detokenize(_sent.split())
+        res = handle_dounble_quote(_sent)
+        if res == -1:
+            print ('unbalanced double quote')
+            print (_sent)
+        else:
+            _sent = res
+        final_sents.append(_sent)
+    sent = ' '.join(final_sents)
+    sent = postprocess_space(sent)
+    return sent