Initial upload: ViBERT-capu ONNX (FP32 + INT8) — converted from dragonSwing/vibert-capu

README.md

@@ -0,0 +1,153 @@
+---
+license: cc-by-sa-4.0
+language:
+- vi
+base_model: dragonSwing/vibert-capu
+tags:
+- punctuation
+- capitalization
+- vietnamese
+- onnx
+- bert
+- vibert
+library_name: onnxruntime
+---
+
+# ViBERT-capu ONNX (FP32 + INT8)
+
+Vietnamese **punctuation restoration + capitalization** model — ONNX Runtime version of [dragonSwing/vibert-capu](https://huggingface.co/dragonSwing/vibert-capu).
+PyTorch dependency removed (~2 GB → ~50 MB onnxruntime).
+
+| Variant | File | Size | Use case |
+|---|---|---|---|
+| FP32 | `vibert-capu.onnx` | 438 MB | Best accuracy, server / web service |
+| INT8 | `vibert-capu.int8.onnx` | 110 MB | Desktop, embedded — dynamic-quantized weights, ~99% of FP32 accuracy |
+
+Architecture: BERT (FPTAI/vibert-base-cased) fine-tuned by [dragonSwing](https://huggingface.co/dragonSwing) on 5.6M OSCAR-2109 samples for the Seq2Labels punctuation+capitalization task (15 GECToR-style edit actions).
+
+## Why ONNX?
+
+| | PyTorch (original) | ONNX Runtime (this repo) |
+|---|---|---|
+| Cold start | ~6 s | ~0.8 s |
+| Runtime deps | torch (~2 GB) | onnxruntime (~50 MB) |
+| Portable build | very heavy | lightweight |
+
+## Quick start
+
+```python
+import numpy as np
+import onnxruntime as ort
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer
+
+local = snapshot_download("welcomyou/vibert-capu-onnx")
+tok = AutoTokenizer.from_pretrained(local)
+sess = ort.InferenceSession(f"{local}/vibert-capu.int8.onnx",
+                            providers=["CPUExecutionProvider"])
+
+text = "hà nội là thủ đô việt nam tôi yêu nó"
+enc = tok(text.split(), is_split_into_words=True, return_tensors="np")
+# input_offsets: index of first subword for each word
+word_ids = enc.word_ids()
+offsets = []
+prev = None
+for i, w in enumerate(word_ids):
+    if w is not None and w != prev:
+        offsets.append(i); prev = w
+input_offsets = np.array([offsets], dtype=np.int64)
+
+logits, detect_logits = sess.run(None, {
+    "input_ids": enc["input_ids"].astype(np.int64),
+    "attention_mask": enc["attention_mask"].astype(np.int64),
+    "token_type_ids": enc["token_type_ids"].astype(np.int64),
+    "input_offsets": input_offsets,
+})
+# logits: (1, num_words, 15)  — 15 GECToR actions
+# detect_logits: (1, num_words, 4) — error detection
+```
+
+## Model I/O
+
+**Inputs** (all `int64`):
+
+| Name | Shape | Description |
+|---|---|---|
+| `input_ids` | `(batch, seq_len)` | BPE token IDs from BertTokenizer |
+| `attention_mask` | `(batch, seq_len)` | 1 = real token, 0 = padding |
+| `token_type_ids` | `(batch, seq_len)` | Segment IDs (always 0) |
+| `input_offsets` | `(batch, num_words)` | Index of first subword for each whitespace-separated word |
+
+**Outputs** (`float32`):
+
+| Name | Shape | Description |
+|---|---|---|
+| `logits` | `(batch, num_words, 15)` | Action probabilities (15 GECToR-style edits) |
+| `detect_logits` | `(batch, num_words, 4)` | Error-detection probabilities |
+
+**15 actions:**
+
+```
+$KEEP                      Giữ nguyên
+$TRANSFORM_CASE_CAPITAL    Viết hoa chữ cái đầu (hà nội → Hà Nội)
+$APPEND_,                  Thêm dấu phẩy
+$APPEND_.                  Thêm dấu chấm
+$TRANSFORM_VERB_VB_VBN     (không dùng cho tiếng Việt)
+$TRANSFORM_CASE_UPPER      Viết hoa toàn bộ (who → WHO)
+$APPEND_:                  Thêm dấu hai chấm
+$APPEND_?                  Thêm dấu hỏi
+$TRANSFORM_VERB_VB_VBC     (không dùng cho tiếng Việt)
+$TRANSFORM_CASE_LOWER      Viết thường
+$TRANSFORM_CASE_CAPITAL_1  Viết hoa ký tự thứ 2
+$TRANSFORM_CASE_UPPER_-1   Viết hoa trừ ký tự cuối
+$MERGE_SPACE               Nối từ
+@@UNKNOWN@@
+@@PADDING@@
+```
+
+## Reproducing the export
+
+```bash
+git clone https://huggingface.co/dragonSwing/vibert-capu
+pip install torch transformers onnxruntime numpy
+
+# Export FP32 + dynamic-quantize INT8 in one step:
+python convert_onnx/export_vibert_onnx.py \
+    --model_dir vibert-capu \
+    --output    vibert-capu.onnx \
+    --opset     14 \
+    --verify
+```
+
+Script: [`convert_onnx/export_vibert_onnx.py`](https://github.com/welcomyou/sherpa-vietnamese-asr/blob/main/convert_onnx/export_vibert_onnx.py).
+
+## Files
+
+```
+config.json                    BERT config (from dragonSwing)
+vocab.txt                      BERT vocabulary (from dragonSwing)
+vocabulary/                    GECToR action labels
+  d_tags.txt
+  labels.txt
+  non_padded_namespaces.txt
+verb-form-vocab.txt            Verb form vocabulary
+vibert-capu.onnx              FP32 ONNX (438 MB)
+vibert-capu.int8.onnx         INT8 ONNX (110 MB)
+configuration_seq2labels.py   Seq2Labels HF config class
+modeling_seq2labels.py        Seq2Labels HF model class (PyTorch reference, not used at runtime)
+gec_model.py                  GECToR inference helpers
+utils.py                      Tokenization helpers
+vocabulary.py                 GECToR Vocabulary class
+```
+
+## Credits & License
+
+- **Original model**: [dragonSwing/vibert-capu](https://huggingface.co/dragonSwing/vibert-capu)
+- **Base BERT**: [FPTAI/vibert-base-cased](https://huggingface.co/FPTAI/vibert-base-cased)
+- **Training data**: [OSCAR-2109](https://huggingface.co/datasets/oscar-corpus/OSCAR-2109) Vietnamese subset (5.6M samples)
+
+License: **CC-BY-SA-4.0** (inherited from dragonSwing/vibert-capu — derivative works must use the same license).
+
+## Used by
+
+- [Sherpa Vietnamese ASR](https://github.com/welcomyou/sherpa-vietnamese-asr) — offline Vietnamese ASR for desktop and web (CPU-only).

config.json

@@ -0,0 +1,26 @@
+{
+  "attention_probs_dropout_prob": 0.1,
+  "bert_vocab_size": 38168,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "label_smoothing": 0.0,
+  "layer_norm_eps": 1e-12,
+  "load_pretrained": false,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_detect_classes": 4,
+  "num_hidden_layers": 12,
+  "num_labels": 15,
+  "pad_token_id": 0,
+  "predictor_dropout": 0.0,
+  "pretrained_name_or_path": "FPTAI/vibert-base-cased",
+  "special_tokens_fix": true,
+  "transformers_version": "4.18.0",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 15
+}

configuration_seq2labels.py

@@ -0,0 +1,62 @@
+from transformers import PretrainedConfig
+
+
+class Seq2LabelsConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`Seq2LabelsModel`]. It is used to
+    instantiate a Seq2Labels model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the Seq2Labels architecture.
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+    Args:
+        vocab_size (`int`, *optional*, defaults to 30522):
+            Vocabulary size of the BERT model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`BertModel`] or [`TFBertModel`].
+        pretrained_name_or_path (`str`, *optional*, defaults to `bert-base-cased`):
+            Pretrained BERT-like model path
+        load_pretrained (`bool`, *optional*, defaults to `False`):
+            Whether to load pretrained model from `pretrained_name_or_path`
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        predictor_dropout (`float`, *optional*):
+            The dropout ratio for the classification head.
+        special_tokens_fix (`bool`, *optional*, defaults to `False`):
+            Whether to add additional tokens to the BERT's embedding layer.
+    Examples:
+    ```python
+    >>> from transformers import BertModel, BertConfig
+    >>> # Initializing a Seq2Labels style configuration
+    >>> configuration = Seq2LabelsConfig()
+    >>> # Initializing a model from the bert-base-uncased style configuration
+    >>> model = Seq2LabelsModel(configuration)
+    >>> # Accessing the model configuration
+    >>> configuration = model.config
+    ```"""
+    model_type = "bert"
+
+    def __init__(
+        self,
+        pretrained_name_or_path="bert-base-cased",
+        vocab_size=15,
+        num_detect_classes=4,
+        load_pretrained=False,
+        initializer_range=0.02,
+        pad_token_id=0,
+        use_cache=True,
+        predictor_dropout=0.0,
+        special_tokens_fix=False,
+        label_smoothing=0.0,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.num_detect_classes = num_detect_classes
+        self.pretrained_name_or_path = pretrained_name_or_path
+        self.load_pretrained = load_pretrained
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        self.predictor_dropout = predictor_dropout
+        self.special_tokens_fix = special_tokens_fix
+        self.label_smoothing = label_smoothing

gec_model.py

@@ -0,0 +1,446 @@
+"""Wrapper of Seq2Labels model. Fixes errors based on model predictions"""
+from collections import defaultdict
+from difflib import SequenceMatcher
+import logging
+import re
+from time import time
+from typing import List, Union
+import warnings
+
+import torch
+from transformers import AutoTokenizer
+from modeling_seq2labels import Seq2LabelsModel
+from vocabulary import Vocabulary
+from utils import PAD, UNK, START_TOKEN, get_target_sent_by_edits
+
+logging.getLogger("werkzeug").setLevel(logging.ERROR)
+logger = logging.getLogger(__file__)
+
+
+class GecBERTModel(torch.nn.Module):
+    def __init__(
+        self,
+        vocab_path=None,
+        model_paths=None,
+        weights=None,
+        device=None,
+        max_len=64,
+        min_len=3,
+        lowercase_tokens=False,
+        log=False,
+        iterations=3,
+        min_error_probability=0.0,
+        confidence=0,
+        resolve_cycles=False,
+        split_chunk=False,
+        chunk_size=48,
+        overlap_size=12,
+        min_words_cut=6,
+        punc_dict={':', ".", ",", "?"},
+    ):
+        r"""
+        Args:
+            vocab_path (`str`):
+                Path to vocabulary directory.
+            model_paths (`List[str]`):
+                List of model paths.
+            weights (`int`, *Optional*, defaults to None):
+                Weights of each model. Only relevant if `is_ensemble is True`.
+            device (`int`, *Optional*, defaults to None):
+                Device to load model. If not set, device will be automatically choose.
+            max_len (`int`, defaults to 64):
+                Max sentence length to be processed (all longer will be truncated).
+            min_len (`int`, defaults to 3):
+                Min sentence length to be processed (all shorted will be returned w/o changes).
+            lowercase_tokens (`bool`, defaults to False):
+                Whether to lowercase tokens.
+            log (`bool`, defaults to False):
+                Whether to enable logging.
+            iterations (`int`, defaults to 3):
+                Max iterations to run during inference.
+            special_tokens_fix (`bool`, defaults to True):
+               Whether to fix problem with [CLS], [SEP] tokens tokenization.
+            min_error_probability (`float`, defaults to `0.0`):
+                Minimum probability for each action to apply.
+            confidence (`float`, defaults to `0.0`):
+                How many probability to add to $KEEP token.
+            split_chunk (`bool`, defaults to False):
+                Whether to split long sentences to multiple segments of `chunk_size`.
+                !Warning: if `chunk_size > max_len`, each segment will be truncate to `max_len`.
+            chunk_size (`int`, defaults to 48):
+                Length of each segment (in words). Only relevant if `split_chunk is True`.
+            overlap_size (`int`, defaults to 12):
+                Overlap size (in words) between two consecutive segments. Only relevant if `split_chunk is True`.
+            min_words_cut (`int`, defaults to 6):
+                Minimun number of words to be cut while merging two consecutive segments.
+                Only relevant if `split_chunk is True`.
+            punc_dict (List[str], defaults to `{':', ".", ",", "?"}`):
+                List of punctuations.
+        """
+        super().__init__()
+        if isinstance(model_paths, str):
+            model_paths = [model_paths]
+        self.model_weights = list(map(float, weights)) if weights else [1] * len(model_paths)
+        self.device = (
+            torch.device("cuda" if torch.cuda.is_available() else "cpu") if device is None else torch.device(device)
+        )
+        self.max_len = max_len
+        self.min_len = min_len
+        self.lowercase_tokens = lowercase_tokens
+        self.min_error_probability = min_error_probability
+        self.vocab = Vocabulary.from_files(vocab_path)
+        self.incorr_index = self.vocab.get_token_index("INCORRECT", "d_tags")
+        self.log = log
+        self.iterations = iterations
+        self.confidence = confidence
+        self.resolve_cycles = resolve_cycles
+        assert (
+            chunk_size > 0 and chunk_size // 2 >= overlap_size
+        ), "Chunk merging required overlap size must be smaller than half of chunk size"
+        self.split_chunk = split_chunk
+        self.chunk_size = chunk_size
+        self.overlap_size = overlap_size
+        self.min_words_cut = min_words_cut
+        self.stride = chunk_size - overlap_size
+        self.punc_dict = punc_dict
+        self.punc_str = '[' + ''.join([f'\{x}' for x in punc_dict]) + ']'
+        # set training parameters and operations
+
+        self.indexers = []
+        self.models = []
+        for model_path in model_paths:
+            model = Seq2LabelsModel.from_pretrained(model_path)
+            config = model.config
+            model_name = config.pretrained_name_or_path
+            special_tokens_fix = config.special_tokens_fix
+            self.indexers.append(self._get_indexer(model_name, special_tokens_fix))
+            model.eval().to(self.device)
+            self.models.append(model)
+
+    def _get_indexer(self, weights_name, special_tokens_fix):
+        tokenizer = AutoTokenizer.from_pretrained(
+            weights_name, do_basic_tokenize=False, do_lower_case=self.lowercase_tokens, model_max_length=1024
+        )
+        # to adjust all tokenizers
+        if hasattr(tokenizer, 'encoder'):
+            tokenizer.vocab = tokenizer.encoder
+        if hasattr(tokenizer, 'sp_model'):
+            tokenizer.vocab = defaultdict(lambda: 1)
+            for i in range(tokenizer.sp_model.get_piece_size()):
+                tokenizer.vocab[tokenizer.sp_model.id_to_piece(i)] = i
+
+        if special_tokens_fix:
+            tokenizer.add_tokens([START_TOKEN])
+            tokenizer.vocab[START_TOKEN] = len(tokenizer) - 1
+        return tokenizer
+    
+    def forward(self, text: Union[str, List[str], List[List[str]]], is_split_into_words=False):
+        # Input type checking for clearer error
+        def _is_valid_text_input(t):
+            if isinstance(t, str):
+                # Strings are fine
+                return True
+            elif isinstance(t, (list, tuple)):
+                # List are fine as long as they are...
+                if len(t) == 0:
+                    # ... empty
+                    return True
+                elif isinstance(t[0], str):
+                    # ... list of strings
+                    return True
+                elif isinstance(t[0], (list, tuple)):
+                    # ... list with an empty list or with a list of strings
+                    return len(t[0]) == 0 or isinstance(t[0][0], str)
+                else:
+                    return False
+            else:
+                return False
+
+        if not _is_valid_text_input(text):
+            raise ValueError(
+                "text input must of type `str` (single example), `List[str]` (batch or single pretokenized example) "
+                "or `List[List[str]]` (batch of pretokenized examples)."
+            )
+        
+        if is_split_into_words:
+            is_batched = isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple))
+        else:
+            is_batched = isinstance(text, (list, tuple))
+            if is_batched:
+                text = [x.split() for x in text]
+            else:
+                text = text.split()
+        
+        if not is_batched:
+            text = [text]
+        
+        return self.handle_batch(text)
+
+    def split_chunks(self, batch):
+        # return batch pairs of indices
+        result = []
+        indices = []
+        for tokens in batch:
+            start = len(result)
+            num_token = len(tokens)
+            if num_token <= self.chunk_size:
+                result.append(tokens)
+            elif num_token > self.chunk_size and num_token < (self.chunk_size * 2 - self.overlap_size):
+                split_idx = (num_token + self.overlap_size + 1) // 2
+                result.append(tokens[:split_idx])
+                result.append(tokens[split_idx - self.overlap_size :])
+            else:
+                for i in range(0, num_token - self.overlap_size, self.stride):
+                    result.append(tokens[i : i + self.chunk_size])
+
+            indices.append((start, len(result)))
+
+        return result, indices
+
+    def check_alnum(self, s):
+        if len(s) < 2:
+            return False
+        return not (s.isalpha() or s.isdigit())
+
+    def apply_chunk_merging(self, tokens, next_tokens):
+        # Return next tokens if current tokens list is empty
+        if not tokens:
+            return next_tokens
+
+        source_token_idx = []
+        target_token_idx = []
+        source_tokens = []
+        target_tokens = []
+        num_keep = self.overlap_size - self.min_words_cut
+        i = 0
+        while len(source_token_idx) < self.overlap_size and -i < len(tokens):
+            i -= 1
+            if tokens[i] not in self.punc_dict:
+                source_token_idx.insert(0, i)
+                source_tokens.insert(0, tokens[i].lower())
+
+        i = 0
+        while len(target_token_idx) < self.overlap_size and i < len(next_tokens):
+            if next_tokens[i] not in self.punc_dict:
+                target_token_idx.append(i)
+                target_tokens.append(next_tokens[i].lower())
+            i += 1
+
+        matcher = SequenceMatcher(None, source_tokens, target_tokens)
+        diffs = list(matcher.get_opcodes())
+
+        for diff in diffs:
+            tag, i1, i2, j1, j2 = diff
+            if tag == "equal":
+                if i1 >= num_keep:
+                    tail_idx = source_token_idx[i1]
+                    head_idx = target_token_idx[j1]
+                    break
+                elif i2 > num_keep:
+                    tail_idx = source_token_idx[num_keep]
+                    head_idx = target_token_idx[j2 - i2 + num_keep]
+                    break
+            elif tag == "delete" and i1 == 0:
+                num_keep += i2 // 2
+
+        tokens = tokens[:tail_idx] + next_tokens[head_idx:]
+        return tokens
+
+    def merge_chunks(self, batch):
+        result = []
+        if len(batch) == 1 or self.overlap_size == 0:
+            for sub_tokens in batch:
+                result.extend(sub_tokens)
+        else:
+            for _, sub_tokens in enumerate(batch):
+                try:
+                    result = self.apply_chunk_merging(result, sub_tokens)
+                except Exception as e:
+                    print(e)
+
+        result = " ".join(result)
+        return result
+
+    def predict(self, batches):
+        t11 = time()
+        predictions = []
+        for batch, model in zip(batches, self.models):
+            batch = batch.to(self.device)
+            with torch.no_grad():
+                prediction = model.forward(**batch)
+            predictions.append(prediction)
+
+        preds, idx, error_probs = self._convert(predictions)
+        t55 = time()
+        if self.log:
+            print(f"Inference time {t55 - t11}")
+        return preds, idx, error_probs
+
+    def get_token_action(self, token, index, prob, sugg_token):
+        """Get lost of suggested actions for token."""
+        # cases when we don't need to do anything
+        if prob < self.min_error_probability or sugg_token in [UNK, PAD, '$KEEP']:
+            return None
+
+        if sugg_token.startswith('$REPLACE_') or sugg_token.startswith('$TRANSFORM_') or sugg_token == '$DELETE':
+            start_pos = index
+            end_pos = index + 1
+        elif sugg_token.startswith("$APPEND_") or sugg_token.startswith("$MERGE_"):
+            start_pos = index + 1
+            end_pos = index + 1
+
+        if sugg_token == "$DELETE":
+            sugg_token_clear = ""
+        elif sugg_token.startswith('$TRANSFORM_') or sugg_token.startswith("$MERGE_"):
+            sugg_token_clear = sugg_token[:]
+        else:
+            sugg_token_clear = sugg_token[sugg_token.index('_') + 1 :]
+
+        return start_pos - 1, end_pos - 1, sugg_token_clear, prob
+
+    def preprocess(self, token_batch):
+        seq_lens = [len(sequence) for sequence in token_batch if sequence]
+        if not seq_lens:
+            return []
+        max_len = min(max(seq_lens), self.max_len)
+        batches = []
+        for indexer in self.indexers:
+            token_batch = [[START_TOKEN] + sequence[:max_len] for sequence in token_batch]
+            batch = indexer(
+                token_batch,
+                return_tensors="pt",
+                padding=True,
+                is_split_into_words=True,
+                truncation=True,
+                add_special_tokens=False,
+            )
+            offset_batch = []
+            for i in range(len(token_batch)):
+                word_ids = batch.word_ids(batch_index=i)
+                offsets = [0]
+                for i in range(1, len(word_ids)):
+                    if word_ids[i] != word_ids[i - 1]:
+                        offsets.append(i)
+                offset_batch.append(torch.LongTensor(offsets))
+
+            batch["input_offsets"] = torch.nn.utils.rnn.pad_sequence(
+                offset_batch, batch_first=True, padding_value=0
+            ).to(torch.long)
+
+            batches.append(batch)
+
+        return batches
+
+    def _convert(self, data):
+        all_class_probs = torch.zeros_like(data[0]['logits'])
+        error_probs = torch.zeros_like(data[0]['max_error_probability'])
+        for output, weight in zip(data, self.model_weights):
+            class_probabilities_labels = torch.softmax(output['logits'], dim=-1)
+            all_class_probs += weight * class_probabilities_labels / sum(self.model_weights)
+            class_probabilities_d = torch.softmax(output['detect_logits'], dim=-1)
+            error_probs_d = class_probabilities_d[:, :, self.incorr_index]
+            incorr_prob = torch.max(error_probs_d, dim=-1)[0]
+            error_probs += weight * incorr_prob / sum(self.model_weights)
+
+        max_vals = torch.max(all_class_probs, dim=-1)
+        probs = max_vals[0].tolist()
+        idx = max_vals[1].tolist()
+        return probs, idx, error_probs.tolist()
+
+    def update_final_batch(self, final_batch, pred_ids, pred_batch, prev_preds_dict):
+        new_pred_ids = []
+        total_updated = 0
+        for i, orig_id in enumerate(pred_ids):
+            orig = final_batch[orig_id]
+            pred = pred_batch[i]
+            prev_preds = prev_preds_dict[orig_id]
+            if orig != pred and pred not in prev_preds:
+                final_batch[orig_id] = pred
+                new_pred_ids.append(orig_id)
+                prev_preds_dict[orig_id].append(pred)
+                total_updated += 1
+            elif orig != pred and pred in prev_preds:
+                # update final batch, but stop iterations
+                final_batch[orig_id] = pred
+                total_updated += 1
+            else:
+                continue
+        return final_batch, new_pred_ids, total_updated
+
+    def postprocess_batch(self, batch, all_probabilities, all_idxs, error_probs):
+        all_results = []
+        noop_index = self.vocab.get_token_index("$KEEP", "labels")
+        for tokens, probabilities, idxs, error_prob in zip(batch, all_probabilities, all_idxs, error_probs):
+            length = min(len(tokens), self.max_len)
+            edits = []
+
+            # skip whole sentences if there no errors
+            if max(idxs) == 0:
+                all_results.append(tokens)
+                continue
+
+            # skip whole sentence if probability of correctness is not high
+            if error_prob < self.min_error_probability:
+                all_results.append(tokens)
+                continue
+
+            for i in range(length + 1):
+                # because of START token
+                if i == 0:
+                    token = START_TOKEN
+                else:
+                    token = tokens[i - 1]
+                # skip if there is no error
+                if idxs[i] == noop_index:
+                    continue
+
+                sugg_token = self.vocab.get_token_from_index(idxs[i], namespace='labels')
+                action = self.get_token_action(token, i, probabilities[i], sugg_token)
+                if not action:
+                    continue
+
+                edits.append(action)
+            all_results.append(get_target_sent_by_edits(tokens, edits))
+        return all_results
+
+    def handle_batch(self, full_batch, merge_punc=True):
+        """
+        Handle batch of requests.
+        """
+        if self.split_chunk:
+            full_batch, indices = self.split_chunks(full_batch)
+        else:
+            indices = None
+        final_batch = full_batch[:]
+        batch_size = len(full_batch)
+        prev_preds_dict = {i: [final_batch[i]] for i in range(len(final_batch))}
+        short_ids = [i for i in range(len(full_batch)) if len(full_batch[i]) < self.min_len]
+        pred_ids = [i for i in range(len(full_batch)) if i not in short_ids]
+        total_updates = 0
+
+        for n_iter in range(self.iterations):
+            orig_batch = [final_batch[i] for i in pred_ids]
+
+            sequences = self.preprocess(orig_batch)
+
+            if not sequences:
+                break
+            probabilities, idxs, error_probs = self.predict(sequences)
+
+            pred_batch = self.postprocess_batch(orig_batch, probabilities, idxs, error_probs)
+            if self.log:
+                print(f"Iteration {n_iter + 1}. Predicted {round(100*len(pred_ids)/batch_size, 1)}% of sentences.")
+
+            final_batch, pred_ids, cnt = self.update_final_batch(final_batch, pred_ids, pred_batch, prev_preds_dict)
+            total_updates += cnt
+
+            if not pred_ids:
+                break
+        if self.split_chunk:
+            final_batch = [self.merge_chunks(final_batch[start:end]) for (start, end) in indices]
+        else:
+            final_batch = [" ".join(x) for x in final_batch]
+        if merge_punc:
+            final_batch = [re.sub(r'\s+(%s)' % self.punc_str, r'\1', x) for x in final_batch]
+
+        return final_batch

modeling_seq2labels.py

@@ -0,0 +1,131 @@
+from typing import Any, Dict, List, Optional, Tuple, Union
+from torch import nn
+from torch.nn import CrossEntropyLoss
+from transformers import AutoConfig, AutoModel, BertPreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+
+import torch
+
+
+def get_range_vector(size: int, device: int) -> torch.Tensor:
+    """
+    Returns a range vector with the desired size, starting at 0. The CUDA implementation
+    is meant to avoid copy data from CPU to GPU.
+    """
+    return torch.arange(0, size, dtype=torch.long, device=device)
+
+
+class Seq2LabelsOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    detect_logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+    max_error_probability: Optional[torch.FloatTensor] = None
+
+
+class Seq2LabelsModel(BertPreTrainedModel):
+
+    _keys_to_ignore_on_load_unexpected = [r"pooler"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.num_detect_classes = config.num_detect_classes
+        self.label_smoothing = config.label_smoothing
+
+        if config.load_pretrained:
+            self.bert = AutoModel.from_pretrained(config.pretrained_name_or_path)
+            bert_config = self.bert.config
+        else:
+            # Uu tien dung config hien tai (da co cac truong BERT architecture)
+            # thay vi goi AutoConfig.from_pretrained() online
+            if hasattr(config, 'hidden_size') and config.hidden_size:
+                from copy import deepcopy
+                bert_config = deepcopy(config)
+                # vocab_size trong config la so labels (15), khong phai BERT tokenizer vocab
+                bert_config.vocab_size = getattr(config, 'bert_vocab_size', 38168)
+            else:
+                bert_config = AutoConfig.from_pretrained(config.pretrained_name_or_path)
+            self.bert = AutoModel.from_config(bert_config)
+
+        if config.special_tokens_fix:
+            try:
+                vocab_size = self.bert.embeddings.word_embeddings.num_embeddings
+            except AttributeError:
+                # reserve more space
+                vocab_size = self.bert.word_embedding.num_embeddings + 5
+            self.bert.resize_token_embeddings(vocab_size + 1)
+
+        predictor_dropout = config.predictor_dropout if config.predictor_dropout is not None else 0.0
+        self.dropout = nn.Dropout(predictor_dropout)
+        self.classifier = nn.Linear(bert_config.hidden_size, config.vocab_size)
+        self.detector = nn.Linear(bert_config.hidden_size, config.num_detect_classes)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        input_offsets: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        token_type_ids: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        d_tags: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.Tensor], Seq2LabelsOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+
+        if input_offsets is not None:
+            # offsets is (batch_size, d1, ..., dn, orig_sequence_length)
+            range_vector = get_range_vector(input_offsets.size(0), device=sequence_output.device).unsqueeze(1)
+            # selected embeddings is also (batch_size * d1 * ... * dn, orig_sequence_length)
+            sequence_output = sequence_output[range_vector, input_offsets]
+
+        logits = self.classifier(self.dropout(sequence_output))
+        logits_d = self.detector(sequence_output)
+
+        loss = None
+        if labels is not None and d_tags is not None:
+            loss_labels_fct = CrossEntropyLoss(label_smoothing=self.label_smoothing)
+            loss_d_fct = CrossEntropyLoss()
+            loss_labels = loss_labels_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            loss_d = loss_d_fct(logits_d.view(-1, self.num_detect_classes), d_tags.view(-1))
+            loss = loss_labels + loss_d
+
+        if not return_dict:
+            output = (logits, logits_d) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return Seq2LabelsOutput(
+            loss=loss,
+            logits=logits,
+            detect_logits=logits_d,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            max_error_probability=torch.ones(logits.size(0), device=logits.device),
+        )

utils.py

@@ -0,0 +1,233 @@
+import os
+from pathlib import Path
+import re
+
+
+VOCAB_DIR = Path(__file__).resolve().parent
+PAD = "@@PADDING@@"
+UNK = "@@UNKNOWN@@"
+START_TOKEN = "$START"
+SEQ_DELIMETERS = {"tokens": " ", "labels": "SEPL|||SEPR", "operations": "SEPL__SEPR"}
+
+
+def get_verb_form_dicts():
+    path_to_dict = os.path.join(VOCAB_DIR, "verb-form-vocab.txt")
+    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
+
+
+ENCODE_VERB_DICT, DECODE_VERB_DICT = get_verb_form_dicts()
+
+
+def get_target_sent_by_edits(source_tokens, edits):
+    target_tokens = source_tokens[:]
+    shift_idx = 0
+    for edit in edits:
+        start, end, label, _ = edit
+        target_pos = start + shift_idx
+        if start < 0:
+            continue
+        elif len(target_tokens) > target_pos:
+            source_token = target_tokens[target_pos]
+        else:
+            source_token = ""
+        if label == "":
+            del target_tokens[target_pos]
+            shift_idx -= 1
+        elif start == end:
+            word = label.replace("$APPEND_", "")
+            # Avoid appending same token twice
+            if (target_pos < len(target_tokens) and target_tokens[target_pos] == word) or (
+                target_pos > 0 and target_tokens[target_pos - 1] == word
+            ):
+                continue
+            target_tokens[target_pos:target_pos] = [word]
+            shift_idx += 1
+        elif label.startswith("$TRANSFORM_"):
+            word = apply_reverse_transformation(source_token, label)
+            if word is None:
+                word = source_token
+            target_tokens[target_pos] = word
+        elif start == end - 1:
+            word = label.replace("$REPLACE_", "")
+            target_tokens[target_pos] = word
+        elif label.startswith("$MERGE_"):
+            target_tokens[target_pos + 1 : target_pos + 1] = [label]
+            shift_idx += 1
+
+    return replace_merge_transforms(target_tokens)
+
+
+def replace_merge_transforms(tokens):
+    if all(not x.startswith("$MERGE_") for x in tokens):
+        return tokens
+    if tokens[0].startswith("$MERGE_"):
+        tokens = tokens[1:]
+    if tokens[-1].startswith("$MERGE_"):
+        tokens = tokens[:-1]
+
+    target_line = " ".join(tokens)
+    target_line = target_line.replace(" $MERGE_HYPHEN ", "-")
+    target_line = target_line.replace(" $MERGE_SPACE ", "")
+    target_line = re.sub(r'([\.\,\?\:]\s+)+', r'\1', target_line)
+    return target_line.split()
+
+
+def convert_using_case(token, smart_action):
+    if not smart_action.startswith("$TRANSFORM_CASE_"):
+        return token
+    if smart_action.endswith("LOWER"):
+        return token.lower()
+    elif smart_action.endswith("UPPER"):
+        return token.upper()
+    elif smart_action.endswith("CAPITAL"):
+        return token.capitalize()
+    elif smart_action.endswith("CAPITAL_1"):
+        return token[0] + token[1:].capitalize()
+    elif smart_action.endswith("UPPER_-1"):
+        return token[:-1].upper() + token[-1]
+    else:
+        return token
+
+
+def convert_using_verb(token, smart_action):
+    key_word = "$TRANSFORM_VERB_"
+    if not smart_action.startswith(key_word):
+        raise Exception(f"Unknown action type {smart_action}")
+    encoding_part = f"{token}_{smart_action[len(key_word):]}"
+    decoded_target_word = decode_verb_form(encoding_part)
+    return decoded_target_word
+
+
+def convert_using_split(token, smart_action):
+    key_word = "$TRANSFORM_SPLIT"
+    if not smart_action.startswith(key_word):
+        raise Exception(f"Unknown action type {smart_action}")
+    target_words = token.split("-")
+    return " ".join(target_words)
+
+
+def convert_using_plural(token, smart_action):
+    if smart_action.endswith("PLURAL"):
+        return token + "s"
+    elif smart_action.endswith("SINGULAR"):
+        return token[:-1]
+    else:
+        raise Exception(f"Unknown action type {smart_action}")
+
+
+def apply_reverse_transformation(source_token, transform):
+    if transform.startswith("$TRANSFORM"):
+        # deal with equal
+        if transform == "$KEEP":
+            return source_token
+        # deal with case
+        if transform.startswith("$TRANSFORM_CASE"):
+            return convert_using_case(source_token, transform)
+        # deal with verb
+        if transform.startswith("$TRANSFORM_VERB"):
+            return convert_using_verb(source_token, transform)
+        # deal with split
+        if transform.startswith("$TRANSFORM_SPLIT"):
+            return convert_using_split(source_token, transform)
+        # deal with single/plural
+        if transform.startswith("$TRANSFORM_AGREEMENT"):
+            return convert_using_plural(source_token, transform)
+        # raise exception if not find correct type
+        raise Exception(f"Unknown action type {transform}")
+    else:
+        return source_token
+
+
+# def read_parallel_lines(fn1, fn2):
+#     lines1 = read_lines(fn1, skip_strip=True)
+#     lines2 = read_lines(fn2, skip_strip=True)
+#     assert len(lines1) == len(lines2)
+#     out_lines1, out_lines2 = [], []
+#     for line1, line2 in zip(lines1, lines2):
+#         if not line1.strip() or not line2.strip():
+#             continue
+#         else:
+#             out_lines1.append(line1)
+#             out_lines2.append(line2)
+#     return out_lines1, out_lines2
+
+
+def read_parallel_lines(fn1, fn2):
+    with open(fn1, encoding='utf-8') as f1, open(fn2, encoding='utf-8') as f2:
+        for line1, line2 in zip(f1, f2):
+            line1 = line1.strip()
+            line2 = line2.strip()
+
+            yield line1, line2
+
+
+def read_lines(fn, skip_strip=False):
+    if not os.path.exists(fn):
+        return []
+    with open(fn, 'r', encoding='utf-8') as f:
+        lines = f.readlines()
+    return [s.strip() for s in lines if s.strip() or skip_strip]
+
+
+def write_lines(fn, lines, mode='w'):
+    if mode == 'w' and os.path.exists(fn):
+        os.remove(fn)
+    with open(fn, encoding='utf-8', mode=mode) as f:
+        f.writelines(['%s\n' % s for s in lines])
+
+
+def decode_verb_form(original):
+    return DECODE_VERB_DICT.get(original)
+
+
+def encode_verb_form(original_word, corrected_word):
+    decoding_request = original_word + "_" + corrected_word
+    decoding_response = ENCODE_VERB_DICT.get(decoding_request, "").strip()
+    if original_word and decoding_response:
+        answer = decoding_response
+    else:
+        answer = None
+    return answer
+
+
+def get_weights_name(transformer_name, lowercase):
+    if transformer_name == 'bert' and lowercase:
+        return 'bert-base-uncased'
+    if transformer_name == 'bert' and not lowercase:
+        return 'bert-base-cased'
+    if transformer_name == 'bert-large' and not lowercase:
+        return 'bert-large-cased'
+    if transformer_name == 'distilbert':
+        if not lowercase:
+            print('Warning! This model was trained only on uncased sentences.')
+        return 'distilbert-base-uncased'
+    if transformer_name == 'albert':
+        if not lowercase:
+            print('Warning! This model was trained only on uncased sentences.')
+        return 'albert-base-v1'
+    if lowercase:
+        print('Warning! This model was trained only on cased sentences.')
+    if transformer_name == 'roberta':
+        return 'roberta-base'
+    if transformer_name == 'roberta-large':
+        return 'roberta-large'
+    if transformer_name == 'gpt2':
+        return 'gpt2'
+    if transformer_name == 'transformerxl':
+        return 'transfo-xl-wt103'
+    if transformer_name == 'xlnet':
+        return 'xlnet-base-cased'
+    if transformer_name == 'xlnet-large':
+        return 'xlnet-large-cased'
+
+    return transformer_name

vibert-capu.int8.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:67278b23502bbd744538e2bd9a6748b61cce171cfba6c671c2e2a46f892166fa
+size 115707731

vibert-capu.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:269a59c50977cef010292b1530a77df3073420ff6be409c1a55eeb77a8444e44
+size 459449409

vocabulary.py

@@ -0,0 +1,277 @@
+import codecs
+from collections import defaultdict
+import logging
+import os
+import re
+from typing import Any, Callable, Dict, Iterable, List, Optional, Set, Union, TYPE_CHECKING
+from filelock import FileLock
+
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_NON_PADDED_NAMESPACES = ("*tags", "*labels")
+DEFAULT_PADDING_TOKEN = "@@PADDING@@"
+DEFAULT_OOV_TOKEN = "@@UNKNOWN@@"
+NAMESPACE_PADDING_FILE = "non_padded_namespaces.txt"
+_NEW_LINE_REGEX = re.compile(r"\n|\r\n")
+
+
+def namespace_match(pattern: str, namespace: str):
+    """
+    Matches a namespace pattern against a namespace string.  For example, `*tags` matches
+    `passage_tags` and `question_tags` and `tokens` matches `tokens` but not
+    `stemmed_tokens`.
+    """
+    if pattern[0] == "*" and namespace.endswith(pattern[1:]):
+        return True
+    elif pattern == namespace:
+        return True
+    return False
+
+
+class _NamespaceDependentDefaultDict(defaultdict):
+    """
+    This is a [defaultdict]
+    (https://docs.python.org/2/library/collections.html#collections.defaultdict) where the
+    default value is dependent on the key that is passed.
+    We use "namespaces" in the :class:`Vocabulary` object to keep track of several different
+    mappings from strings to integers, so that we have a consistent API for mapping words, tags,
+    labels, characters, or whatever else you want, into integers.  The issue is that some of those
+    namespaces (words and characters) should have integers reserved for padding and
+    out-of-vocabulary tokens, while others (labels and tags) shouldn't.  This class allows you to
+    specify filters on the namespace (the key used in the `defaultdict`), and use different
+    default values depending on whether the namespace passes the filter.
+    To do filtering, we take a set of `non_padded_namespaces`.  This is a set of strings
+    that are either matched exactly against the keys, or treated as suffixes, if the
+    string starts with `*`.  In other words, if `*tags` is in `non_padded_namespaces` then
+    `passage_tags`, `question_tags`, etc. (anything that ends with `tags`) will have the
+    `non_padded` default value.
+    # Parameters
+    non_padded_namespaces : `Iterable[str]`
+        A set / list / tuple of strings describing which namespaces are not padded.  If a namespace
+        (key) is missing from this dictionary, we will use :func:`namespace_match` to see whether
+        the namespace should be padded.  If the given namespace matches any of the strings in this
+        list, we will use `non_padded_function` to initialize the value for that namespace, and
+        we will use `padded_function` otherwise.
+    padded_function : `Callable[[], Any]`
+        A zero-argument function to call to initialize a value for a namespace that `should` be
+        padded.
+    non_padded_function : `Callable[[], Any]`
+        A zero-argument function to call to initialize a value for a namespace that should `not` be
+        padded.
+    """
+
+    def __init__(
+        self,
+        non_padded_namespaces: Iterable[str],
+        padded_function: Callable[[], Any],
+        non_padded_function: Callable[[], Any],
+    ) -> None:
+        self._non_padded_namespaces = set(non_padded_namespaces)
+        self._padded_function = padded_function
+        self._non_padded_function = non_padded_function
+        super().__init__()
+
+    def add_non_padded_namespaces(self, non_padded_namespaces: Set[str]):
+        # add non_padded_namespaces which weren't already present
+        self._non_padded_namespaces.update(non_padded_namespaces)
+
+
+class _TokenToIndexDefaultDict(_NamespaceDependentDefaultDict):
+    def __init__(self, non_padded_namespaces: Set[str], padding_token: str, oov_token: str) -> None:
+        super().__init__(non_padded_namespaces, lambda: {padding_token: 0, oov_token: 1}, lambda: {})
+
+
+class _IndexToTokenDefaultDict(_NamespaceDependentDefaultDict):
+    def __init__(self, non_padded_namespaces: Set[str], padding_token: str, oov_token: str) -> None:
+        super().__init__(non_padded_namespaces, lambda: {0: padding_token, 1: oov_token}, lambda: {})
+
+
+class Vocabulary:
+    def __init__(
+        self,
+        counter: Dict[str, Dict[str, int]] = None,
+        min_count: Dict[str, int] = None,
+        max_vocab_size: Union[int, Dict[str, int]] = None,
+        non_padded_namespaces: Iterable[str] = DEFAULT_NON_PADDED_NAMESPACES,
+        pretrained_files: Optional[Dict[str, str]] = None,
+        only_include_pretrained_words: bool = False,
+        tokens_to_add: Dict[str, List[str]] = None,
+        min_pretrained_embeddings: Dict[str, int] = None,
+        padding_token: Optional[str] = DEFAULT_PADDING_TOKEN,
+        oov_token: Optional[str] = DEFAULT_OOV_TOKEN,
+    ) -> None:
+        self._padding_token = padding_token if padding_token is not None else DEFAULT_PADDING_TOKEN
+        self._oov_token = oov_token if oov_token is not None else DEFAULT_OOV_TOKEN
+
+        self._non_padded_namespaces = set(non_padded_namespaces)
+
+        self._token_to_index = _TokenToIndexDefaultDict(
+            self._non_padded_namespaces, self._padding_token, self._oov_token
+        )
+        self._index_to_token = _IndexToTokenDefaultDict(
+            self._non_padded_namespaces, self._padding_token, self._oov_token
+        )
+
+    @classmethod
+    def from_files(
+        cls,
+        directory: Union[str, os.PathLike],
+        padding_token: Optional[str] = DEFAULT_PADDING_TOKEN,
+        oov_token: Optional[str] = DEFAULT_OOV_TOKEN,
+    ) -> "Vocabulary":
+        """
+        Loads a `Vocabulary` that was serialized either using `save_to_files` or inside
+        a model archive file.
+        # Parameters
+        directory : `str`
+            The directory or archive file containing the serialized vocabulary.
+        """
+        logger.info("Loading token dictionary from %s.", directory)
+        padding_token = padding_token if padding_token is not None else DEFAULT_PADDING_TOKEN
+        oov_token = oov_token if oov_token is not None else DEFAULT_OOV_TOKEN
+
+        if not os.path.isdir(directory):
+            raise ValueError(f"{directory} not exist")
+
+        # We use a lock file to avoid race conditions where multiple processes
+        # might be reading/writing from/to the same vocab files at once.
+        with FileLock(os.path.join(directory, ".lock")):
+            with codecs.open(os.path.join(directory, NAMESPACE_PADDING_FILE), "r", "utf-8") as namespace_file:
+                non_padded_namespaces = [namespace_str.strip() for namespace_str in namespace_file]
+
+            vocab = cls(
+                non_padded_namespaces=non_padded_namespaces,
+                padding_token=padding_token,
+                oov_token=oov_token,
+            )
+
+            # Check every file in the directory.
+            for namespace_filename in os.listdir(directory):
+                if namespace_filename == NAMESPACE_PADDING_FILE:
+                    continue
+                if namespace_filename.startswith("."):
+                    continue
+                namespace = namespace_filename.replace(".txt", "")
+                if any(namespace_match(pattern, namespace) for pattern in non_padded_namespaces):
+                    is_padded = False
+                else:
+                    is_padded = True
+                filename = os.path.join(directory, namespace_filename)
+                vocab.set_from_file(filename, is_padded, namespace=namespace, oov_token=oov_token)
+
+        return vocab
+
+    @classmethod
+    def empty(cls) -> "Vocabulary":
+        """
+        This method returns a bare vocabulary instantiated with `cls()` (so, `Vocabulary()` if you
+        haven't made a subclass of this object).  The only reason to call `Vocabulary.empty()`
+        instead of `Vocabulary()` is if you are instantiating this object from a config file.  We
+        register this constructor with the key "empty", so if you know that you don't need to
+        compute a vocabulary (either because you're loading a pre-trained model from an archive
+        file, you're using a pre-trained transformer that has its own vocabulary, or something
+        else), you can use this to avoid having the default vocabulary construction code iterate
+        through the data.
+        """
+        return cls()
+
+    def set_from_file(
+        self,
+        filename: str,
+        is_padded: bool = True,
+        oov_token: str = DEFAULT_OOV_TOKEN,
+        namespace: str = "tokens",
+    ):
+        """
+        If you already have a vocabulary file for a trained model somewhere, and you really want to
+        use that vocabulary file instead of just setting the vocabulary from a dataset, for
+        whatever reason, you can do that with this method.  You must specify the namespace to use,
+        and we assume that you want to use padding and OOV tokens for this.
+        # Parameters
+        filename : `str`
+            The file containing the vocabulary to load.  It should be formatted as one token per
+            line, with nothing else in the line.  The index we assign to the token is the line
+            number in the file (1-indexed if `is_padded`, 0-indexed otherwise).  Note that this
+            file should contain the OOV token string!
+        is_padded : `bool`, optional (default=`True`)
+            Is this vocabulary padded?  For token / word / character vocabularies, this should be
+            `True`; while for tag or label vocabularies, this should typically be `False`.  If
+            `True`, we add a padding token with index 0, and we enforce that the `oov_token` is
+            present in the file.
+        oov_token : `str`, optional (default=`DEFAULT_OOV_TOKEN`)
+            What token does this vocabulary use to represent out-of-vocabulary characters?  This
+            must show up as a line in the vocabulary file.  When we find it, we replace
+            `oov_token` with `self._oov_token`, because we only use one OOV token across
+            namespaces.
+        namespace : `str`, optional (default=`"tokens"`)
+            What namespace should we overwrite with this vocab file?
+        """
+        if is_padded:
+            self._token_to_index[namespace] = {self._padding_token: 0}
+            self._index_to_token[namespace] = {0: self._padding_token}
+        else:
+            self._token_to_index[namespace] = {}
+            self._index_to_token[namespace] = {}
+        with codecs.open(filename, "r", "utf-8") as input_file:
+            lines = _NEW_LINE_REGEX.split(input_file.read())
+            # Be flexible about having final newline or not
+            if lines and lines[-1] == "":
+                lines = lines[:-1]
+            for i, line in enumerate(lines):
+                index = i + 1 if is_padded else i
+                token = line.replace("@@NEWLINE@@", "\n")
+                if token == oov_token:
+                    token = self._oov_token
+                self._token_to_index[namespace][token] = index
+                self._index_to_token[namespace][index] = token
+        if is_padded:
+            assert self._oov_token in self._token_to_index[namespace], "OOV token not found!"
+
+    def add_token_to_namespace(self, token: str, namespace: str = "tokens") -> int:
+        """
+        Adds `token` to the index, if it is not already present.  Either way, we return the index of
+        the token.
+        """
+        if not isinstance(token, str):
+            raise ValueError(
+                "Vocabulary tokens must be strings, or saving and loading will break."
+                "  Got %s (with type %s)" % (repr(token), type(token))
+            )
+        if token not in self._token_to_index[namespace]:
+            index = len(self._token_to_index[namespace])
+            self._token_to_index[namespace][token] = index
+            self._index_to_token[namespace][index] = token
+            return index
+        else:
+            return self._token_to_index[namespace][token]
+
+    def add_tokens_to_namespace(self, tokens: List[str], namespace: str = "tokens") -> List[int]:
+        """
+        Adds `tokens` to the index, if they are not already present.  Either way, we return the
+        indices of the tokens in the order that they were given.
+        """
+        return [self.add_token_to_namespace(token, namespace) for token in tokens]
+
+    def get_token_index(self, token: str, namespace: str = "tokens") -> int:
+        try:
+            return self._token_to_index[namespace][token]
+        except KeyError:
+            try:
+                return self._token_to_index[namespace][self._oov_token]
+            except KeyError:
+                logger.error("Namespace: %s", namespace)
+                logger.error("Token: %s", token)
+                raise KeyError(
+                    f"'{token}' not found in vocab namespace '{namespace}', and namespace "
+                    f"does not contain the default OOV token ('{self._oov_token}')"
+                )
+
+    def get_token_from_index(self, index: int, namespace: str = "tokens") -> str:
+        return self._index_to_token[namespace][index]
+
+    def get_vocab_size(self, namespace: str = "tokens") -> int:
+        return len(self._token_to_index[namespace])
+
+    def get_namespaces(self) -> Set[str]:
+        return set(self._index_to_token.keys())

vocabulary/d_tags.txt

@@ -0,0 +1,4 @@
+CORRECT
+INCORRECT
+@@UNKNOWN@@
+@@PADDING@@

vocabulary/labels.txt

@@ -0,0 +1,15 @@
+$KEEP
+$TRANSFORM_CASE_CAPITAL
+$APPEND_,
+$APPEND_.
+$TRANSFORM_VERB_VB_VBN
+$TRANSFORM_CASE_UPPER
+$APPEND_:
+$APPEND_?
+$TRANSFORM_VERB_VB_VBC
+$TRANSFORM_CASE_LOWER
+$TRANSFORM_CASE_CAPITAL_1
+$TRANSFORM_CASE_UPPER_-1
+$MERGE_SPACE
+@@UNKNOWN@@
+@@PADDING@@

vocabulary/non_padded_namespaces.txt

@@ -0,0 +1,2 @@
+*tags
+*labels