Upload 9 files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+config/ms_yahei.ttf filter=lfs diff=lfs merge=lfs -text

config.json

@@ -0,0 +1,35 @@
+{
+  "_name_or_path": "iioSnail/NamBert-for-csc",
+  "architectures": [
+    "NamBertForCSC"
+  ],
+  "attention_probs_dropout_prob": 0.1,
+  "auto_map": {
+    "AutoModel": "csc_model.NamBertForCSC"
+  },
+  "classifier_dropout": null,
+  "directionality": "bidi",
+  "gradient_checkpointing": false,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.1,
+  "hidden_size": 768,
+  "initializer_range": 0.02,
+  "intermediate_size": 3072,
+  "layer_norm_eps": 1e-12,
+  "max_position_embeddings": 512,
+  "model_type": "bert",
+  "num_attention_heads": 12,
+  "num_hidden_layers": 12,
+  "pad_token_id": 0,
+  "pooler_fc_size": 768,
+  "pooler_num_attention_heads": 12,
+  "pooler_num_fc_layers": 3,
+  "pooler_size_per_head": 128,
+  "pooler_type": "first_token_transform",
+  "position_embedding_type": "absolute",
+  "torch_dtype": "float32",
+  "transformers_version": "4.33.2",
+  "type_vocab_size": 2,
+  "use_cache": true,
+  "vocab_size": 21128
+}

config/ms_yahei.ttf

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

csc_model.py

@@ -0,0 +1,152 @@
+import os
+import shutil
+import time
+from pathlib import Path
+
+import torch
+from huggingface_hub import hf_hub_download
+from torch import nn
+from torchvision.models import resnet18
+from transformers import BertPreTrainedModel, BertForMaskedLM
+from transformers.activations import GELUActivation
+from transformers.modeling_outputs import MaskedLMOutput
+
+cache_path = Path(os.path.abspath(__file__)).parent
+
+
+def download_file(filename: str, path: Path):
+    if os.path.exists(cache_path / filename):
+        return
+
+    if os.path.exists(path / filename):
+        shutil.copyfile(path / filename, cache_path / filename)
+        return
+
+    hf_hub_download(
+        "iioSnail/NamBert-for-csc",
+        filename,
+        local_dir=cache_path,
+    )
+    time.sleep(0.2)
+
+
+class NamBertForCSC(BertPreTrainedModel):
+
+    def __init__(self, config):
+        super(NamBertForCSC, self).__init__(config)
+
+        self.bert = BertForMaskedLM(config).bert
+        self.token_forget_gate = nn.Linear(768, 768, bias=False)
+
+        self.pinyin_feature_size = 6
+        self.pinyin_embeddings = PinyinManualEmbeddings()
+
+        self.glyph_feature_size = 56
+        self.glyph_embeddings = GlyphDenseEmbedding()
+
+        self.cls = BertOnlyMLMHead(768 + self.pinyin_feature_size + self.glyph_feature_size, config.vocab_size,
+                                   layer_num=1)
+
+    def forward(self, input_ids, attention_mask, token_type_ids, pinyin_ids, images, **kwargs):
+        batch_size = input_ids.size(0)
+
+        del kwargs['offset_mapping']
+
+        bert_outputs = self.bert(input_ids=input_ids,
+                                 attention_mask=attention_mask,
+                                 token_type_ids=token_type_ids,
+                                 **kwargs)
+
+        token_embeddings = self.bert.embeddings(input_ids)
+        token_embeddings = token_embeddings * self.token_forget_gate(token_embeddings).sigmoid()
+        bert_outputs.last_hidden_state += token_embeddings
+
+        pinyin_embeddings = self.pinyin_embeddings(pinyin_ids)
+        pinyin_embeddings = pinyin_embeddings.view(batch_size, -1, self.pinyin_feature_size)
+
+        glyph_embeddings = self.glyph_embeddings(images)
+        glyph_embeddings = glyph_embeddings.view(batch_size, -1, self.glyph_feature_size)
+
+        hidden_states = torch.concat([bert_outputs.last_hidden_state,
+                                      pinyin_embeddings,
+                                      glyph_embeddings], dim=-1)
+
+        logits = self.cls(hidden_states)
+
+        return MaskedLMOutput(
+            logits=logits,
+            hidden_states=bert_outputs.hidden_states,
+            attentions=bert_outputs.attentions,
+        )
+
+
+class PinyinManualEmbeddings(nn.Module):
+
+    def __init__(self):
+        super(PinyinManualEmbeddings, self).__init__()
+        self.pinyin_feature_size = 6
+        self.embedding_layer = nn.Linear(6, 6, bias=True)
+
+    def forward(self, inputs):
+        fill = self.pinyin_feature_size - inputs.size(1)
+        if fill > 0:
+            inputs = torch.concat([inputs, torch.zeros((len(inputs), fill), device=inputs.device)], dim=1).long()
+        inputs = self.embedding_layer(inputs.float())
+        return inputs
+
+
+class GlyphDenseEmbedding(nn.Module):
+
+    def __init__(self, font_size=32):
+        super(GlyphDenseEmbedding, self).__init__()
+        self.font_size = font_size
+        self.embeddings = nn.Sequential(
+            nn.Linear(1024, 512),
+            nn.ReLU(),
+            nn.Dropout(0.15),
+            nn.Linear(512, 256),
+            nn.ReLU(),
+            nn.Dropout(0.15),
+            nn.Linear(256, 56),
+            nn.Tanh()
+        )
+
+    def forward(self, images):
+        batch_size = len(images)
+        images = images.view(batch_size, -1) / 255.
+        return self.embeddings(images)
+
+
+class BertOnlyMLMHead(nn.Module):
+    def __init__(self, hidden_size, vocab_size, activation='gelu', layer_num=1):
+        super().__init__()
+        self.decoder = nn.Linear(hidden_size, vocab_size, bias=False)
+        self.bias = nn.Parameter(torch.zeros(vocab_size))
+        self.decoder.bias = self.bias
+
+        self.activation = None
+        if activation == 'gelu':
+            self.activation = GELUActivation()
+        elif activation == 'tanh':
+            self.activation = nn.Tanh()
+        elif activation == 'sigmoid':
+            self.activation = nn.Sigmoid()
+        else:
+            raise Exception("Please add activation function here.")
+
+        self.heads = []
+
+        for i in range(layer_num):
+            self.heads.append(nn.Sequential(
+                nn.Linear(hidden_size, hidden_size),
+                self.activation,
+                nn.LayerNorm(hidden_size, eps=1e-12, elementwise_affine=True),
+            ))
+
+        self.predictions = nn.Sequential(
+            *self.heads,
+            self.decoder
+        )
+
+    def forward(self, sequence_output: torch.Tensor) -> torch.Tensor:
+        return self.predictions(sequence_output)

csc_tokenizer.py

@@ -0,0 +1,202 @@
+import os
+import shutil
+import time
+from pathlib import Path
+from typing import List, Union, Optional
+
+import numpy as np
+import pypinyin
+import torch
+from PIL import ImageFont
+from huggingface_hub import hf_hub_download
+from torch import NoneType
+from torch.nn.utils.rnn import pad_sequence
+from transformers.tokenization_utils_base import TruncationStrategy
+from transformers.utils import PaddingStrategy
+from transformers.utils.generic import TensorType
+
+try:
+    from tokenizers import BertWordPieceTokenizer
+except:
+    from tokenizers.implementations import BertWordPieceTokenizer
+
+from transformers import BertTokenizerFast, BatchEncoding
+
+cache_path = Path(os.path.abspath(__file__)).parent
+
+
+def download_file(filename: str, path: Path):
+    if os.path.exists(cache_path / filename):
+        return
+
+    if os.path.exists(path / filename):
+        shutil.copyfile(path / filename, cache_path / filename)
+        return
+
+    hf_hub_download(
+        "iioSnail/NamBert-for-csc",
+        filename,
+        local_dir=cache_path
+    )
+    time.sleep(0.2)
+
+
+class NamBertTokenizer(BertTokenizerFast):
+
+    def __init__(self, **kwargs):
+        super(NamBertTokenizer, self).__init__(**kwargs)
+
+        self.path = Path(kwargs['name_or_path'])
+        vocab_file = cache_path / 'vocab.txt'
+        config_path = cache_path / 'config'
+        if not os.path.exists(config_path):
+            os.makedirs(config_path)
+
+        self.max_length = 20480
+        self.font_size = 32
+
+        download_file('vocab.txt', self.path)
+        download_file('config/ms_yahei.ttf', self.path)
+
+        self.tokenizer = BertWordPieceTokenizer(str(vocab_file))
+
+        font = ImageFont.truetype(str(cache_path / 'config' / "ms_yahei.ttf"), size=self.font_size)
+        vocab = self.tokenizer.get_vocab().items()
+        self.input_helper = InputHelper(font, vocab)
+
+    def __call__(self,
+                 text: Union[str, List[str], List[List[str]]] = None,
+                 text_pair: Union[str, List[str], List[List[str]], NoneType] = None,
+                 text_target: Union[str, List[str], List[List[str]]] = None,
+                 text_pair_target: Union[str, List[str], List[List[str]], NoneType] = None,
+                 add_special_tokens: bool = True,
+                 padding: Union[bool, str, PaddingStrategy] = False,
+                 truncation: Union[bool, str, TruncationStrategy] = None,
+                 max_length: Optional[int] = None,
+                 stride: int = 0,
+                 is_split_into_words: bool = False,
+                 pad_to_multiple_of: Optional[int] = None,
+                 return_tensors: Union[str, TensorType, NoneType] = None,
+                 return_token_type_ids: Optional[bool] = None,
+                 return_attention_mask: Optional[bool] = None,
+                 return_overflowing_tokens: bool = False, return_special_tokens_mask: bool = False,
+                 return_offsets_mapping: bool = False,
+                 return_length: bool = False,
+                 verbose: bool = True, **kwargs) -> BatchEncoding:
+        encoding = super(NamBertTokenizer, self).__call__(
+            text=text,
+            text_pair=text_pair,
+            text_target=text_target,
+            text_pair_target=text_pair_target,
+            add_special_tokens=add_special_tokens,
+            padding=padding,
+            truncation=truncation,
+            max_length=max_length,
+            stride=stride,
+            is_split_into_words=is_split_into_words,
+            pad_to_multiple_of=pad_to_multiple_of,
+            return_tensors='pt',
+            return_token_type_ids=return_token_type_ids,
+            return_attention_mask=return_attention_mask,
+            return_overflowing_tokens=return_overflowing_tokens,
+            return_offsets_mapping=True,
+            return_length=return_length,
+            verbose=verbose,
+        )
+
+        input_ids = encoding.input_ids
+        encoding['pinyin_ids'] = self.input_helper.convert_tokens_to_pinyin_embeddings(input_ids.view(-1))
+        encoding['images'] = self.input_helper.convert_tokens_to_images(input_ids.view(-1))
+
+        return encoding
+
+    def restore_ids(self, target_ids, input_ids):
+        for i in range(len(target_ids)):
+            for j in range(len(target_ids[i])):
+                if target_ids[i][j] == 1 or target_ids[i][j] == 0:
+                    target_ids[i][j] = input_ids[i][j]
+
+        return target_ids
+
+
+class InputHelper:
+
+    def __init__(self, font, vocab):
+        self.font = font
+        self.vocab = vocab
+
+        self.pinyin_embedding_cache = None
+        self._init_pinyin_embedding_cache()
+
+        self.token_images_cache = None
+        self._init_token_images_cache()
+
+    def _init_pinyin_embedding_cache(self):
+        self.pinyin_embedding_cache = {}
+        for token, id in self.vocab:
+            self.pinyin_embedding_cache[id] = convert_char_to_pinyin(token)
+
+    def _init_token_images_cache(self):
+        self.token_images_cache = {}
+        for token, id in self.vocab:
+            self.token_images_cache[id] = convert_char_to_image(self.font, token, 32)
+
+    def convert_tokens_to_pinyin_embeddings(self, input_ids):
+        input_pinyins = []
+        for i, input_id in enumerate(input_ids):
+            input_pinyins.append(self.pinyin_embedding_cache.get(input_id.item(), torch.LongTensor([0])))
+
+        return pad_sequence(input_pinyins, batch_first=True)
+
+    def convert_tokens_to_images(self, input_ids):
+        images = []
+        for i, input_id in enumerate(input_ids):
+            images.append(self.token_images_cache.get(input_id.item(), torch.zeros(32, 32)))
+        return torch.stack(images)
+
+
+def convert_char_to_pinyin(character, size=-1, tone=False):
+    if not is_chinese(character):
+        return torch.LongTensor([0] * max(size, 1))
+
+    if tone:
+        pinyin = pypinyin.pinyin(character, style=pypinyin.TONE3)[0][0]
+    else:
+        pinyin = pypinyin.pinyin(character, style=pypinyin.NORMAL)[0][0]
+
+    if not tone:
+        embeddings = torch.tensor([ord(letter) - 96 for letter in pinyin])
+    else:
+        embeddings = []
+        for letter in pinyin:
+            if letter.isnumeric():
+                embeddings.append(int(letter) + 27)
+            else:
+                embeddings.append(ord(letter) - 96)
+        embeddings = torch.tensor(embeddings)
+
+    if size > len(embeddings):
+        padding = torch.zeros(size - len(embeddings))
+        embeddings = torch.concat([embeddings, padding])
+
+    return embeddings
+
+
+def convert_char_to_image(font, character, font_size=32):
+    image = font.getmask(character)
+    image = np.asarray(image).astype(np.float32).reshape(image.size[::-1])
+
+    image = image[:font_size, :font_size]
+
+    if image.size != (font_size, font_size):
+        back_image = np.zeros((font_size, font_size)).astype(np.float32)
+        offset0 = (font_size - image.shape[0]) // 2
+        offset1 = (font_size - image.shape[1]) // 2
+        back_image[offset0:offset0 + image.shape[0], offset1:offset1 + image.shape[1]] = image
+        image = back_image
+
+    return torch.tensor(image)
+
+
+def is_chinese(uchar):
+    return '\u4e00' <= uchar <= '\u9fa5'

pytorch_model.bin

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

special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "[CLS]",
+  "mask_token": "[MASK]",
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "unk_token": "[UNK]"
+}

tokenizer_config.json

@@ -0,0 +1,19 @@
+{
+  "auto_map": {
+    "AutoTokenizer": [
+      "csc_tokenizer.NamBertTokenizer",
+      null
+    ]
+  },
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "do_lower_case": true,
+  "mask_token": "[MASK]",
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "strip_accents": null,
+  "tokenize_chinese_chars": true,
+  "tokenizer_class": "NamBertTokenizer",
+  "unk_token": "[UNK]"
+}