add handler

gector/__init__.py

@@ -0,0 +1,22 @@
+from .modeling import GECToR
+from .configuration import GECToRConfig
+from .dataset import load_dataset, GECToRDataset
+from .predict import predict, load_verb_dict
+from .predict_verbose import predict_verbose
+from .vocab import (
+    build_vocab,
+    load_vocab_from_config,
+    load_vocab_from_official
+)
+__all__ = [
+    'GECToR',
+    'GECToRConfig',
+    'load_dataset',
+    'GECToRDataset',
+    'predict',
+    'load_verb_dict',
+    'predict_verbose',
+    'build_vocab',
+    'load_vocab_from_config',
+    'load_vocab_from_official'
+]

gector/configuration.py

@@ -0,0 +1,38 @@
+import os
+import json
+from transformers import PretrainedConfig
+class GECToRConfig(PretrainedConfig):
+    def __init__(
+        self,
+        model_id: str = 'bert-base-cased',
+        p_dropout: float=0,
+        label_pad_token: str='<PAD>',
+        label_oov_token: str='<OOV>',
+        d_pad_token: str='<PAD>',
+        keep_label: str='$KEEP',
+        correct_label: str='$CORRECT',
+        incorrect_label: str='$INCORRECT',
+        label_smoothing: float=0.0,
+        has_add_pooling_layer: bool=True,
+        initializer_range: float=0.02,
+        **kwards
+    ):
+        super().__init__(**kwards)
+        self.d_label2id = {
+            "$CORRECT": 0,
+            "$INCORRECT": 1,
+            "<PAD>": 2
+        }
+        self.d_id2label = {v: k for k, v in self.d_label2id.items()}
+        self.d_num_labels = len(self.d_label2id)
+        self.model_id = model_id
+        self.p_dropout = p_dropout
+        self.label_pad_token = label_pad_token
+        self.label_oov_token = label_oov_token
+        self.d_pad_token = d_pad_token
+        self.keep_label = keep_label
+        self.correct_label = correct_label
+        self.incorrect_label = incorrect_label
+        self.label_smoothing = label_smoothing
+        self.has_add_pooling_layer = has_add_pooling_layer
+        self.initializer_range = initializer_range

gector/dataset.py

@@ -0,0 +1,164 @@
+from typing import List, Tuple
+from collections import Counter
+import torch
+from tqdm import tqdm
+import os
+from transformers import PreTrainedTokenizer
+
+class GECToRDataset:
+    def __init__(
+        self,
+        srcs: List[str],
+        d_labels: List[List[int]]=None,
+        labels: List[List[int]]=None,
+        word_masks: List[List[int]]=None,
+        tokenizer: PreTrainedTokenizer=None,
+        max_length:int=128
+    ):
+        self.tokenizer = tokenizer
+        self.srcs = srcs
+        self.d_labels = d_labels
+        self.labels = labels
+        self.word_masks = word_masks
+        self.max_length = max_length
+        self.label2id = None
+        self.d_label2id = None
+        
+    def __len__(self):
+        return len(self.srcs)
+
+    def __getitem__(self, idx):
+        src = self.srcs[idx]
+        d_labels = self.d_labels[idx]
+        labels = self.labels[idx]
+        wmask = self.word_masks[idx]
+        encode = self.tokenizer(
+            src,
+            return_tensors='pt',
+            max_length=self.max_length,
+            padding='max_length',
+            truncation=True,
+            is_split_into_words=True
+        )
+        return {
+            'input_ids': encode['input_ids'].squeeze(),
+            'attention_mask': encode['attention_mask'].squeeze(),
+            'd_labels': torch.tensor(d_labels).squeeze(),
+            'labels': torch.tensor(labels).squeeze(),
+            'word_masks': torch.tensor(wmask).squeeze()
+        }
+
+    def append_vocab(self, label2id, d_label2id):
+        self.label2id = label2id
+        self.d_label2id = d_label2id
+        for i in range(len(self.labels)):
+            self.labels[i] = [self.label2id.get(l, self.label2id['<OOV>']) for l in self.labels[i]]
+            self.d_labels[i] = [self.d_label2id[l] for l in self.d_labels[i]]
+    
+    def get_labels_freq(self, exluded_labels: List[str] = []):
+        assert(self.labels is not None and self.d_labels is not None)
+        flatten_labels = [ll for l in self.labels for ll in l if ll not in exluded_labels]
+        flatten_d_labels = [ll for l in self.d_labels for ll in l if ll not in exluded_labels]
+        return Counter(flatten_labels), Counter(flatten_d_labels)
+
+def align_labels_to_subwords(
+    srcs: List[str],
+    word_labels: List[List[str]],
+    tokenizer: PreTrainedTokenizer,
+    batch_size: int=100000,
+    max_length: int=128,
+    keep_label: str='$KEEP',
+    pad_token: str='<PAD>',
+    correct_label: str='$CORRECT',
+    incorrect_label: str='$INCORRECT'
+):
+    itr = list(range(0, len(srcs), batch_size))
+    subword_labels = []
+    subword_d_labels = []
+    word_masks = []
+    for i in tqdm(itr):
+        encode = tokenizer(
+            srcs[i:i+batch_size],
+            max_length=max_length,
+            return_tensors='pt',
+            padding='max_length',
+            truncation=True,
+            is_split_into_words=True
+        )
+        for i, wlabels in enumerate(word_labels[i:i+batch_size]):
+            d_labels = []
+            labels = []
+            wmask = []
+            word_ids = encode.word_ids(i)
+            previous_word_idx = None
+            for word_idx in word_ids:
+                if word_idx is None:
+                    labels.append(pad_token)
+                    d_labels.append(pad_token)
+                    wmask.append(0)
+                elif word_idx != previous_word_idx:
+                    l = wlabels[word_idx]
+                    labels.append(l)
+                    wmask.append(1)
+                    if l != keep_label:
+                        d_labels.append(incorrect_label)
+                    else:
+                        d_labels.append(correct_label)
+                else:
+                    labels.append(pad_token)
+                    d_labels.append(pad_token)
+                    wmask.append(0)
+                previous_word_idx = word_idx
+            subword_d_labels.append(d_labels)
+            subword_labels.append(labels)
+            word_masks.append(wmask)
+    return subword_d_labels, subword_labels, word_masks
+        
+def load_gector_format(
+    input_file: str,
+    delimeter: str='SEPL|||SEPR',
+    additional_delimeter: str='SEPL__SEPR'
+):  
+    srcs = []
+    word_level_labels = []  # the size will be (#sents, seq_length) if not get_interactive_tags,
+                                # (#iteration, #sents, seq_length) if get_interactive_tags
+    with open(input_file) as f:
+        for line in f:
+            src = [x.split(delimeter)[0] for x in line.split()]
+            labels = [x.split(delimeter)[1] for x in line.split()]
+            # Use only first tags. E.g. $REPLACE_meSEPL__SEPR$APPEND_too → $REPLACE_me
+            labels = [l.split(additional_delimeter)[0] for l in labels]
+            srcs.append(src)
+            word_level_labels.append(labels)
+    return srcs, word_level_labels
+
+def load_dataset(
+    input_file: str,
+    tokenizer: PreTrainedTokenizer,
+    delimeter: str='SEPL|||SEPR',
+    additional_delimeter: str='SEPL__SEPR',
+    batch_size: int=50000, # avoid too heavy computation in the tokenization
+    max_length: int=128
+):
+    srcs, word_level_labels = load_gector_format(
+        input_file,
+        delimeter=delimeter,
+        additional_delimeter=additional_delimeter
+    )
+    d_labels, labels, word_masks = align_labels_to_subwords(
+        srcs,
+        word_level_labels,
+        tokenizer=tokenizer,
+        batch_size=batch_size,
+        max_length=max_length
+    )
+    return GECToRDataset(
+        srcs=srcs,
+        d_labels=d_labels,
+        labels=labels,
+        word_masks=word_masks,
+        tokenizer=tokenizer,
+        max_length=max_length
+    )
+    
+    

gector/modeling.py

@@ -0,0 +1,200 @@
+from transformers import AutoModel, AutoTokenizer, AutoConfig, PreTrainedModel
+import torch
+import torch.nn.functional as F
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss
+from dataclasses import dataclass
+from .configuration import GECToRConfig
+from typing import List, Union, Optional, Tuple
+import os
+import json
+from huggingface_hub import snapshot_download, ModelCard
+
+@dataclass
+class GECToROutput:
+    loss: torch.Tensor = None
+    loss_d: torch.Tensor = None
+    loss_labels: torch.Tensor = None
+    logits_d: torch.Tensor = None
+    logits_labels: torch.Tensor = None
+    accuracy: torch.Tensor = None
+    accuracy_d: torch.Tensor = None
+
+@dataclass
+class GECToRPredictionOutput:
+    probability_labels: torch.Tensor = None
+    probability_d: torch.Tensor = None
+    pred_labels: List[List[str]] = None
+    pred_label_ids: torch.Tensor = None
+    max_error_probability: torch.Tensor = None
+
+class GECToR(PreTrainedModel):
+    config_class = GECToRConfig
+    def __init__(
+        self,
+        config: GECToRConfig
+    ):
+        super().__init__(config)
+        self.config = config
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config.model_id
+        )
+        if self.config.has_add_pooling_layer:
+            self.bert = AutoModel.from_pretrained(
+                self.config.model_id,
+                add_pooling_layer=False
+            )
+        else:
+            self.bert = AutoModel.from_pretrained(
+                self.config.model_id
+            )
+        # +1 is for $START token
+        self.bert.resize_token_embeddings(self.bert.config.vocab_size + 1)
+        self.label_proj_layer = nn.Linear(
+            self.bert.config.hidden_size,
+            self.config.num_labels - 1
+        )  # -1 is for <PAD>
+        self.d_proj_layer = nn.Linear(
+            self.bert.config.hidden_size,
+            self.config.d_num_labels - 1
+        )
+        self.dropout = nn.Dropout(self.config.p_dropout)
+        self.loss_fn = CrossEntropyLoss(
+            label_smoothing=self.config.label_smoothing
+        )
+        
+        self.post_init()
+        self.tune_bert(False)
+
+    def init_weight(self) -> None:
+        self._init_weights(self.label_proj_layer)
+        self._init_weights(self.d_proj_layer)
+
+    def _init_weights(self, module) -> None:
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(
+                mean=0.0,
+                std=self.config.initializer_range
+            )
+            if module.bias is not None:
+                module.bias.data.zero_()
+        return
+
+    def tune_bert(self, tune=True):
+        # If tune=False, only classifier layers will be tuned.
+        for param in self.bert.parameters():
+            param.requires_grad = tune
+        return
+    
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        d_labels: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        word_masks: Optional[torch.Tensor] = None,
+    ) -> GECToROutput:
+        bert_logits = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        ).last_hidden_state
+        logits_d = self.d_proj_layer(bert_logits)
+        logits_labels = self.label_proj_layer(self.dropout(bert_logits))
+        loss_d, loss_labels, loss = None, None, None
+        accuracy, accuracy_d = None, None
+        if d_labels is not None and labels is not None:
+            pad_id = self.config.label2id[self.config.label_pad_token]
+            # -100 is the default ignore_idx of CrossEntropyLoss
+            labels[labels == pad_id] = -100
+            d_labels[labels == -100] = -100
+            loss_d = self.loss_fn(
+                logits_d.view(-1, self.config.d_num_labels - 1),  # -1 for <PAD>
+                d_labels.view(-1)
+            )
+            loss_labels = self.loss_fn(
+                logits_labels.view(-1, self.config.num_labels - 1),
+                labels.view(-1)
+            )
+            loss = loss_d + loss_labels
+
+            pred_labels = torch.argmax(logits_labels, dim=-1)
+            accuracy = torch.sum(
+                (labels == pred_labels) * word_masks
+            ) / torch.sum(word_masks)
+            pred_d = torch.argmax(logits_d, dim=-1)
+            accuracy_d = torch.sum(
+                (d_labels == pred_d) * word_masks
+            ) / torch.sum(word_masks)
+
+        return GECToROutput(
+            loss=loss,
+            loss_d=loss_d,
+            loss_labels=loss_labels,
+            logits_d=logits_d,
+            logits_labels=logits_labels,
+            accuracy=accuracy,
+            accuracy_d=accuracy_d
+        )
+
+    def predict(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor,
+        word_masks: torch.Tensor,
+        keep_confidence: float=0,
+        min_error_prob: float=0
+    ): 
+        with torch.no_grad():
+            outputs = self.forward(
+                input_ids,
+                attention_mask
+            )
+            probability_labels = F.softmax(outputs.logits_labels, dim=-1)
+            probability_d = F.softmax(outputs.logits_d, dim=-1)
+
+            # Get actual labels considering inference parameters.
+            keep_index = self.config.label2id[self.config.keep_label]
+            probability_labels[:, :, keep_index] += keep_confidence
+            incor_idx = self.config.d_label2id[self.config.incorrect_label]
+            probability_d = probability_d[:, :, incor_idx]
+            max_error_probability = torch.max(probability_d * word_masks, dim=-1)[0]
+            probability_labels[max_error_probability < min_error_prob, :, keep_index] \
+                                                                        = float('inf')
+            pred_label_ids = torch.argmax(probability_labels, dim=-1)
+
+            def convert_ids_to_labels(ids, id2label):
+                labels = []
+                for id in ids.tolist():
+                    labels.append(id2label[id])
+                return labels
+
+            pred_labels = []
+            for ids in pred_label_ids:
+                labels = convert_ids_to_labels(
+                    ids,
+                    self.config.id2label
+                )
+                pred_labels.append(labels)
+
+        return GECToRPredictionOutput(
+            probability_labels=probability_labels,
+            probability_d=probability_d,
+            pred_labels=pred_labels,
+            pred_label_ids=pred_label_ids,
+            max_error_probability=max_error_probability
+        )

gector/predict.py

@@ -0,0 +1,232 @@
+import torch
+import os
+from tqdm import tqdm
+from .modeling import GECToR
+from transformers import PreTrainedTokenizer
+from typing import List
+
+def load_verb_dict(verb_file: str):
+    path_to_dict = os.path.join(verb_file)
+    encode, decode = {}, {}
+    with open(path_to_dict, encoding="utf-8") as f:
+        for line in f:
+            words, tags = line.split(":")
+            word1, word2 = words.split("_")
+            tag1, tag2 = tags.split("_")
+            decode_key = f"{word1}_{tag1}_{tag2.strip()}"
+            if decode_key not in decode:
+                encode[words] = tags
+                decode[decode_key] = word2
+    return encode, decode
+
+def edit_src_by_tags(
+    srcs: List[List[str]],
+    pred_labels: List[List[str]],
+    encode: dict,
+    decode: dict
+) -> List[str]:
+    edited_srcs = []
+    for tokens, labels in zip(srcs, pred_labels):
+        edited_tokens = []
+        for t, l, in zip(tokens, labels):
+            n_token = process_token(t, l, encode, decode)
+            if n_token == None:
+                n_token = t
+            edited_tokens += n_token.split(' ')
+        if len(tokens) > len(labels):
+            omitted_tokens = tokens[len(labels):]
+            edited_tokens += omitted_tokens
+        temp_str = ' '.join(edited_tokens) \
+            .replace(' $MERGE_HYPHEN ', '-') \
+            .replace(' $MERGE_SPACE ', '') \
+            .replace(' $DELETE', '') \
+            .replace('$DELETE ', '')
+        edited_srcs.append(temp_str.split(' '))
+    return edited_srcs
+
+def process_token(
+    token: str,
+    label: str,
+    encode: dict,
+    decode: dict
+) -> str:
+    if '$APPEND_' in label:
+        return token + ' ' + label.replace('$APPEND_', '')
+    elif token == '$START':
+        # [unused1] token cannot be replaced with another token and cannot be deleted.
+        return token
+    elif label in ['<PAD>', '<OOV>', '$KEEP']:
+        return token
+    elif '$APPEND_' in label:
+        return token + ' ' + label.replace('$APPEND_', '')
+    elif '$TRANSFORM_' in label:
+        return g_transform_processer(token, label, encode, decode)
+    elif '$REPLACE_' in label:
+        return label.replace('$REPLACE_', '')
+    elif label == '$DELETE':
+        return label
+    elif '$MERGE_' in label:
+        return token + ' ' + label
+    else:
+        return token
+    
+def g_transform_processer(
+    token: str,
+    label: str,
+    encode: dict,
+    decode: dict
+) -> str:
+    # Case related
+    if label == '$TRANSFORM_CASE_LOWER':
+        return token.lower()
+    elif label == '$TRANSFORM_CASE_UPPER':
+        return token.upper()
+    elif label == '$TRANSFORM_CASE_CAPITAL':
+        return token.capitalize()
+    elif label == '$TRANSFORM_CASE_CAPITAL_1':
+        if len(token) <= 1:
+            return token
+        return token[0] + token[1:].capitalize()
+    elif label == '$TRANSFORM_AGREEMENT_PLURAL':
+        return token + 's'
+    elif label == '$TRANSFORM_AGREEMENT_SINGULAR':
+        return token[:-1]
+    elif label == '$TRANSFORM_SPLIT_HYPHEN':
+        return ' '.join(token.split('-'))
+    else:
+        encoding_part = f"{token}_{label[len('$TRANSFORM_VERB_'):]}"
+        decoded_target_word = decode.get(encoding_part)
+        return decoded_target_word
+
+def get_word_masks_from_word_ids(
+    word_ids: List[List[int]],
+    n: int
+):
+    word_masks = []
+    for i in range(n):
+        previous_id = 0
+        mask = []
+        for _id in word_ids(i):
+            if _id is None:
+                mask.append(0)
+            elif previous_id != _id:
+                mask.append(1)
+            else:
+                mask.append(0)
+            previous_id = _id
+        word_masks.append(mask)
+    return word_masks
+
+def _predict(
+    model: GECToR,
+    tokenizer: PreTrainedTokenizer,
+    srcs: List[str],
+    keep_confidence: float=0,
+    min_error_prob: float=0,
+    batch_size: int=128
+):
+    itr = list(range(0, len(srcs), batch_size))
+    pred_labels = []
+    no_corrections = []
+    for i in tqdm(itr):
+        batch = tokenizer(
+            srcs[i:i+batch_size],
+            return_tensors='pt',
+            max_length=model.config.max_length,
+            padding='max_length',
+            truncation=True,
+            is_split_into_words=True
+        )
+        batch['word_masks'] = torch.tensor(
+            get_word_masks_from_word_ids(
+                batch.word_ids,
+                batch['input_ids'].size(0)
+            )
+        )
+        word_ids = batch.word_ids
+        if torch.cuda.is_available():
+            batch = {k:v.cuda() for k,v in batch.items()}
+        outputs = model.predict(
+            batch['input_ids'],
+            batch['attention_mask'],
+            batch['word_masks'],
+            keep_confidence,
+            min_error_prob
+        )
+        # Align subword-level label to word-level label
+        for i in range(len(outputs.pred_labels)):
+            no_correct = True
+            labels = []
+            previous_word_idx = None
+            for j, idx in enumerate(word_ids(i)):
+                if idx is None:
+                    continue
+                if idx != previous_word_idx:
+                    labels.append(outputs.pred_labels[i][j])
+                    if outputs.pred_label_ids[i][j] > 2:
+                        no_correct = False
+                previous_word_idx = idx
+            # print(no_correct, labels)
+            pred_labels.append(labels)
+            no_corrections.append(no_correct)
+    # print(pred_labels)
+    return pred_labels, no_corrections
+
+def predict(
+    model: GECToR,
+    tokenizer: PreTrainedTokenizer,
+    srcs: List[str],
+    encode: dict,
+    decode: dict,
+    keep_confidence: float=0,
+    min_error_prob: float=0,
+    batch_size: int=128,
+    n_iteration: int=5
+) -> List[str]:
+    srcs = [['$START'] + src.split(' ') for src in srcs]
+    final_edited_sents = ['-1'] * len(srcs)
+    to_be_processed = srcs
+    original_sent_idx = list(range(0, len(srcs)))
+    for itr in range(n_iteration):
+        print(f'Iteratoin {itr}. the number of to_be_processed: {len(to_be_processed)}')
+        pred_labels, no_corrections = _predict(
+            model,
+            tokenizer,
+            to_be_processed,
+            keep_confidence,
+            min_error_prob,
+            batch_size
+        )
+        current_srcs = []
+        current_pred_labels = []
+        current_orig_idx = []
+        for i, yes in enumerate(no_corrections):
+            if yes: # there's no corrections?
+                final_edited_sents[original_sent_idx[i]] = ' '.join(to_be_processed[i]).replace('$START ', '')
+            else:
+                current_srcs.append(to_be_processed[i])
+                current_pred_labels.append(pred_labels[i])
+                current_orig_idx.append(original_sent_idx[i])
+        if current_srcs == []:
+            # Correcting for all sentences is completed.
+            break
+        # if itr > 2:
+        #     for l in current_pred_labels:
+        #         print(l)
+        edited_srcs = edit_src_by_tags(
+            current_srcs,
+            current_pred_labels,
+            encode,
+            decode
+        )
+        to_be_processed = edited_srcs
+        original_sent_idx = current_orig_idx
+        
+        # print(f'=== Iteration {itr} ===')
+        # print('\n'.join(final_edited_sents))
+        # print(to_be_processed)
+        # print(have_corrections)
+    for i in range(len(to_be_processed)):
+        final_edited_sents[original_sent_idx[i]] = ' '.join(to_be_processed[i]).replace('$START ', '')
+    assert('-1' not in final_edited_sents)
+    return final_edited_sents

gector/predict_verbose.py

@@ -0,0 +1,83 @@
+import torch
+import os
+from tqdm import tqdm
+from .modeling import GECToR
+from transformers import PreTrainedTokenizer
+from typing import List, Dict
+from .predict import (
+    edit_src_by_tags,
+    _predict
+)
+
+def predict_verbose(
+    model: GECToR,
+    tokenizer: PreTrainedTokenizer,
+    srcs: List[str],
+    encode: dict,
+    decode: dict,
+    keep_confidence: float=0,
+    min_error_prob: float=0,
+    batch_size: int=128,
+    n_iteration: int=5
+) -> List[str]:
+    srcs = [['$START'] + src.split(' ') for src in srcs]
+    final_edited_sents = ['-1'] * len(srcs)
+    to_be_processed = srcs
+    original_sent_idx = list(range(0, len(srcs)))
+    iteration_log: List[List[Dict]] = []  # [send_id][iteration_id]['src' or 'tags']
+    iteration_log = []
+    # Initialize iteration logs. 
+    for i, src in enumerate(srcs):
+        iteration_log.append([{
+            'src': src,
+            'tag': None
+        }])
+    for itr in range(n_iteration):
+        print(f'Iteratoin {itr}. the number of to_be_processed: {len(to_be_processed)}')
+        pred_labels, no_corrections = _predict(
+            model,
+            tokenizer,
+            to_be_processed,
+            keep_confidence,
+            min_error_prob,
+            batch_size
+        )
+        current_srcs = []
+        current_pred_labels = []
+        current_orig_idx = []
+        for i, yes in enumerate(no_corrections):
+            if yes: # there's no corrections?
+                final_edited_sents[original_sent_idx[i]] = ' '.join(to_be_processed[i]).replace('$START ', '')
+            else:
+                current_srcs.append(to_be_processed[i])
+                current_pred_labels.append(pred_labels[i])
+                current_orig_idx.append(original_sent_idx[i])
+        if current_srcs == []:
+            # Correcting for all sentences is completed.
+            break
+        edited_srcs = edit_src_by_tags(
+            current_srcs,
+            current_pred_labels,
+            encode,
+            decode
+        )
+        # Register the information during iteration.
+        # edited_src will be the src of the next iteration.
+        for i, orig_id in enumerate(current_orig_idx):
+            iteration_log[orig_id][itr]['tag'] = current_pred_labels[i]
+            iteration_log[orig_id].append({
+                'src': edited_srcs[i],
+                'tag': None
+            })
+        
+        to_be_processed = edited_srcs
+        original_sent_idx = current_orig_idx
+        
+        # print(f'=== Iteration {itr} ===')
+        # print('\n'.join(final_edited_sents))
+        # print(to_be_processed)
+        # print(have_corrections)
+    for i in range(len(to_be_processed)):
+        final_edited_sents[original_sent_idx[i]] = ' '.join(to_be_processed[i]).replace('$START ', '')
+    assert('-1' not in final_edited_sents)
+    return final_edited_sents, iteration_log

gector/vocab.py

@@ -0,0 +1,48 @@
+from .configuration import GECToRConfig
+from .dataset import GECToRDataset
+import os
+
+def build_vocab(
+    train_dataset: GECToRDataset,
+    n_max_labels: int=5000,
+    n_max_d_labels: int=2
+):
+    label2id = {'<OOV>':0, '$KEEP':1}
+    d_label2id = {'$CORRECT':0, '$INCORRECT':1, '<PAD>':2}
+    freq_labels, _ = train_dataset.get_labels_freq(
+        exluded_labels=['<PAD>'] + list(label2id.keys())
+    )
+
+    def get_high_freq(freq: dict, n_max: int):
+        descending_freq = sorted(
+            freq.items(), key=lambda x:x[1], reverse=True
+        )
+        high_freq = [x[0] for x in descending_freq][:n_max]
+        if len(high_freq) < n_max:
+            print(f'Warning: the size of the vocablary: {len(high_freq)} is less than n_max: {n_max}.')
+        return high_freq
+    
+    high_freq_labels = get_high_freq(freq_labels, n_max_labels-2)
+    for i, x in enumerate(high_freq_labels):
+        label2id[x] = i + 2
+    label2id['<PAD>'] = len(label2id)
+    return label2id, d_label2id
+
+def load_vocab_from_config(config_file: str):
+    config = GECToRConfig.from_pretrained(config_file, not_dir=True)
+    return config.label2id, config.d_label2id
+
+def load_vocab_from_official(dir):
+    vocab_path = os.path.join(dir, 'labels.txt')
+    vocab = open(vocab_path).read().replace('@@PADDING@@', '').replace('@@UNKNOWN@@', '').rstrip().split('\n')
+    # vocab_d = open(dir + 'd_tags.txt').read().rstrip().replace('@@PADDING@@', '<PAD>').replace('@@UNKNOWN@@', '<OOV>').split('\n')
+    label2id = {'<OOV>':0, '$KEEP':1}
+    d_label2id = {'$CORRECT':0, '$INCORRECT':1, '<PAD>':2}
+    idx = len(label2id)
+    for v in vocab:
+        if v not in label2id:
+            label2id[v] = idx
+            idx += 1
+    label2id['<PAD>'] = idx
+    return label2id, d_label2id
+        

handler.py

@@ -0,0 +1,45 @@
+from typing import Dict, List, Any
+from transformers import AutoTokenizer
+from gector import GECToR, predict, load_verb_dict
+
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.model = GECToR.from_pretrained(path)
+        self.tokenizer = AutoTokenizer.from_pretrained(path)
+        self.encode, self.decode = load_verb_dict("data/verb-form-vocab.txt")
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Process the input data and return the predicted results.
+
+        Args:
+            data (Dict[str, Any]): The input data dictionary containing the following keys:
+                - "inputs" (List[str]): A list of input strings to be processed.
+                - "n_iterations" (int, optional): The number of iterations for prediction. Defaults to 5.
+                - "batch_size" (int, optional): The batch size for prediction. Defaults to 2.
+                - "keep_confidence" (float, optional): The confidence threshold for keeping predictions. Defaults to 0.0.
+                - "min_error_prob" (float, optional): The minimum error probability for keeping predictions. Defaults to 0.0.
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing the predicted results for each input string.
+        """
+        srcs = data["inputs"]
+
+        # Extract optional parameters from data, with defaults
+        n_iterations = data.get("n_iterations", 5)
+        batch_size = data.get("batch_size", 2)
+        keep_confidence = data.get("keep_confidence", 0.0)
+        min_error_prob = data.get("min_error_prob", 0.0)
+
+        return predict(
+            model=self.model,
+            tokenizer=self.tokenizer,
+            srcs=srcs,
+            encode=self.encode,
+            decode=self.decode,
+            keep_confidence=keep_confidence,
+            min_error_prob=min_error_prob,
+            n_iteration=n_iterations,
+            batch_size=batch_size,
+        )

requirements.txt

@@ -0,0 +1,27 @@
+accelerate==0.27.0
+certifi==2024.2.2
+charset-normalizer==3.3.2
+filelock==3.13.1
+fsspec==2024.2.0
+huggingface-hub==0.20.3
+idna==3.6
+Jinja2==3.1.3
+Levenshtein==0.24.0
+MarkupSafe==2.1.5
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.4
+packaging==23.2
+psutil==5.9.8
+PyYAML==6.0.1
+rapidfuzz==3.6.1
+regex==2023.12.25
+requests==2.31.0
+safetensors==0.4.2
+sympy==1.12
+tokenizers==0.15.1
+torch==2.2.0
+tqdm==4.66.2
+transformers==4.37.2
+typing_extensions==4.9.0
+urllib3==2.2.0