upload files

README.md

@@ -0,0 +1,26 @@
+```python
+import torch
+from transformers import AutoModel, AutoTokenizer
+
+model_path = 'CjangCjengh/NomBert-hn2qn-v0.1'
+device = 'cuda'
+
+model = AutoModel.from_pretrained(model_path, torch_dtype='auto', trust_remote_code=True).eval().to(device)
+tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
+
+with torch.inference_mode():
+    output_text, output_probs = model.parse_nom_text(tokenizer, ['仍調𬖉𧡊㐌𤴬疸𢚸'])
+    print(output_text[0])
+    # những điều trông thấy đã đau đớn lòng
+    print(output_probs[0])
+    # [
+    # {'char': '仍', 'candidates': [('những', 0.5237383842468262), ('nhưng', 0.475042462348938), ('dưng', 0.0008663760963827372), ('nhang', 0.00022805406479164958), ('dừng', 8.42325171106495e-05), ('nhẵng', 1.6380783563363366e-05), ('nhùng', 1.5950208762660623e-05), ('nhửng', 3.0440487535088323e-06), ('nhăng', 2.9528700906666927e-06), ('nhẳng', 1.0688020211091498e-06), ('nhừng', 5.84112399337755e-07), ('nhâng', 5.119333650327462e-07)]},
+    # {'char': '調', 'candidates': [('điều', 0.8831620812416077), ('đều', 0.11558306217193604), ('điệu', 0.0012446790933609009), ('dìu', 8.889981472748332e-06), ('điu', 7.615183221787447e-07), ('đìu', 5.942594043517602e-07)]},
+    # {'char': '𬖉', 'candidates': [('trông', 1.0)]},
+    # {'char': '𧡊', 'candidates': [('thấy', 1.0)]},
+    # {'char': '㐌', 'candidates': [('đã', 0.9998464584350586), ('dã', 0.00014108473260421306), ('đà', 1.2395633348205592e-05)]},
+    # {'char': '𤴬', 'candidates': [('đau', 0.9999825954437256), ('đáu', 1.744620021781884e-05)]},
+    # {'char': '疸', 'candidates': [('đớn', 0.9998302459716797), ('đơn', 0.00014517175441142172), ('đảm', 2.457975824654568e-05)]},
+    # {'char': '𢚸', 'candidates': [('lòng', 1.0)]}
+    # ]
+```

configuration_nombert.py

@@ -0,0 +1,10 @@
+from transformers import BertConfig
+
+
+class NomBertConfig(BertConfig):
+    def __init__(self, unk_id=0, id_start=1, output_vocab_size=7430, lm_head_dict={}, **kwargs):
+        super().__init__(**kwargs)
+        self.unk_id = unk_id
+        self.id_start = id_start
+        self.output_vocab_size = output_vocab_size
+        self.lm_head_dict = lm_head_dict

model.safetensors

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

modeling_nombert.py

@@ -0,0 +1,234 @@
+import re
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import BertPreTrainedModel, BertModel
+from .configuration_nombert import NomBertConfig
+
+
+class NomBertModel(BertPreTrainedModel):
+    config_class = NomBertConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.bert = BertModel(config)
+        self.max_position_embeddings = config.max_position_embeddings
+        self.lm_head_dict = config.lm_head_dict
+        self.registered_token_ids = list(map(int, config.lm_head_dict.keys()))
+        self.lm_head = nn.Embedding(config.output_vocab_size, config.hidden_size)
+
+    def forward(self, input_ids, labels=None, attention_mask=None):
+        outputs = self.bert(input_ids, attention_mask)
+        hidden_states = outputs.last_hidden_state
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(input_ids, dtype=torch.long)
+
+        registered_token_ids_tensor = torch.tensor(
+            self.registered_token_ids, 
+            device=input_ids.device
+        )
+        valid_token_mask = torch.isin(input_ids, registered_token_ids_tensor)
+        valid_mask = valid_token_mask & attention_mask.bool()
+
+        loss = torch.tensor(0.0, device=input_ids.device, requires_grad=True)
+        
+        for token_id_str in self.lm_head_dict.keys():
+            token_id = int(token_id_str)
+            mask = (input_ids == token_id) & valid_mask
+            selected_hidden = hidden_states[mask]
+            selected_labels = labels[mask] if labels is not None else None
+            
+            if selected_hidden.size(0) == 0:
+                continue
+
+            lm_head_ids = self.lm_head_dict[token_id_str]
+            lm_head_ids_tensor = torch.tensor(lm_head_ids, device=input_ids.device)
+            lm_head = self.lm_head(lm_head_ids_tensor)
+            logits = torch.matmul(selected_hidden, lm_head.T)
+            
+            if labels is not None:
+                loss = loss + F.cross_entropy(
+                    logits, 
+                    selected_labels, 
+                    ignore_index=-100
+                )
+        
+        return {'loss': loss} if labels is not None else outputs
+
+    def parse_nom_text(self, tokenizer, texts, post_normalize=True, batch_size=None):
+        max_length = self.max_position_embeddings
+        segments_info = []
+        for text_idx, text in enumerate(texts):
+            segments = [text[i:i+max_length] for i in range(0, len(text), max_length)]
+            for seg_idx, seg in enumerate(segments):
+                segments_info.append((text_idx, seg_idx, seg))
+        
+        all_segments = [seg for _, _, seg in segments_info]
+        all_pred_chars = []
+        all_pred_probs = []
+
+        if batch_size is None:
+            batch_size = len(texts)
+        for i in range(0, len(all_segments), batch_size):
+            batch_segments = all_segments[i:i+batch_size]
+            batch_pred_chars, batch_pred_probs = self._parse_nom_text_batch(tokenizer, batch_segments)
+            all_pred_chars.extend(batch_pred_chars)
+            all_pred_probs.extend(batch_pred_probs)
+
+        text_results = {}
+        for text_idx in range(len(texts)):
+            text_results[text_idx] = {'chars': [], 'probs': []}
+        
+        for (text_idx, seg_idx, _), pred_chars, pred_probs in zip(segments_info, all_pred_chars, all_pred_probs):
+            text_results[text_idx]['chars'].append((seg_idx, pred_chars))
+            text_results[text_idx]['probs'].append((seg_idx, pred_probs))
+
+        output_texts = []
+        all_outputs_probs = []
+        for text_idx in range(len(texts)):
+            sorted_chars = sorted(text_results[text_idx]['chars'], key=lambda x: x[0])
+            sorted_probs = sorted(text_results[text_idx]['probs'], key=lambda x: x[0])
+            
+            merged_chars = []
+            merged_probs = []
+            for seg_idx, chars in sorted_chars:
+                merged_chars.extend(chars)
+            for seg_idx, probs in sorted_probs:
+                merged_probs.extend(probs)
+
+            output_text = ''
+            for i, (char, processed) in enumerate(merged_chars):
+                output_text += char
+                if i < len(merged_chars)-1 and (processed or merged_chars[i+1][1]):
+                    output_text += ' '
+
+            if post_normalize:
+                output_text = self.post_normalize(output_text)
+            output_texts.append(output_text)
+            all_outputs_probs.append(merged_probs)
+
+        return output_texts, all_outputs_probs
+
+    def _parse_nom_text_batch(self, tokenizer, segments):
+        encoded = tokenizer.batch_encode_plus(
+            segments,
+            add_special_tokens=False,
+            padding=True,
+            return_tensors='pt',
+            truncation=True,
+            max_length=self.max_position_embeddings
+        )
+        input_ids = encoded['input_ids'].to(self.device)
+        attention_mask = encoded['attention_mask'].to(self.device)
+
+        batch_size = len(segments)
+        id_to_options_ids = list(tokenizer.id_to_options.keys())
+        id_to_options_tensor = torch.tensor(id_to_options_ids, device=self.device)
+        registered_ids = torch.tensor(self.registered_token_ids, device=self.device)
+        valid_mask = (
+            torch.isin(input_ids, registered_ids) & 
+            attention_mask.bool()
+        )
+
+        pred_chars = [[(c, False) for c in seg] for seg in segments]
+        pred_probs = [[] for _ in range(batch_size)]
+
+        if valid_mask.any():
+            outputs = self.bert(input_ids, attention_mask=attention_mask)
+            hidden_states = outputs.last_hidden_state
+            
+            batch_indices, seq_indices = torch.where(valid_mask)
+            token_ids = input_ids[batch_indices, seq_indices]
+            hidden_vecs = hidden_states[batch_indices, seq_indices]
+
+            for token_id_str in self.lm_head_dict:
+                token_id = int(token_id_str)
+                token_mask = (token_ids == token_id)
+                if not token_mask.any():
+                    continue
+
+                token_hidden = hidden_vecs[token_mask]
+                token_batch = batch_indices[token_mask]
+                token_seq = seq_indices[token_mask]
+
+                lm_head_ids = self.lm_head_dict[token_id_str]
+                lm_head_ids_tensor = torch.tensor(lm_head_ids, device=input_ids.device)
+                lm_head = self.lm_head(lm_head_ids_tensor)
+                logits = torch.matmul(token_hidden, lm_head.T)
+                probs = F.softmax(logits, dim=-1)
+                preds = torch.argmax(logits, dim=-1)
+
+                for i, (b, s) in enumerate(zip(token_batch.tolist(), token_seq.tolist())):
+                    options = tokenizer.id_to_options[token_id]
+                    char = options[preds[i].item()]
+
+                    pred_chars[b][s] = (char, True)
+
+                    candidates = sorted(
+                        [(opt, probs[i][j].item()) for j, opt in enumerate(options)],
+                        key=lambda x: x[1], reverse=True
+                    )
+
+                    if s >= len(pred_probs[b]):
+                        pred_probs[b].extend([{}] * (s - len(pred_probs[b]) + 1))
+
+                    pred_probs[b][s] = {
+                        'char': segments[b][s],
+                        'candidates': candidates
+                    }
+
+        single_option_mask = (
+            attention_mask.bool() &
+            torch.isin(input_ids, id_to_options_tensor) &
+            ~torch.isin(input_ids, registered_ids)
+        )
+        batch_indices_single, seq_indices_single = torch.where(single_option_mask)
+
+        for b, s in zip(batch_indices_single.tolist(), seq_indices_single.tolist()):
+            token_id = input_ids[b, s].item()
+            options = tokenizer.id_to_options[token_id]
+
+            pred_chars[b][s] = (options[0], True)
+
+            if s >= len(pred_probs[b]):
+                pred_probs[b].extend([{}] * (s - len(pred_probs[b]) + 1))
+
+            pred_probs[b][s] = {
+                'char': segments[b][s],
+                'candidates': [(options[0], 1.0)]
+            }
+
+        for b in range(batch_size):
+            seg_len = len(segments[b])
+            pred_chars[b] = pred_chars[b][:seg_len]
+            for s in range(seg_len):
+                if s < len(pred_probs[b]) and pred_probs[b][s]:
+                    continue
+                char = segments[b][s]
+                if s >= input_ids.shape[1]:
+                    token_id = 0
+                else:
+                    token_id = input_ids[b, s].item()
+
+                candidates = [(char, 1.0)]
+                if token_id != 0 and token_id in tokenizer.id_to_options:
+                    options = tokenizer.id_to_options[token_id]
+                    if len(options) == 1:
+                        candidates = [(options[0], 1.0)]
+
+                if s >= len(pred_probs[b]):
+                    pred_probs[b].extend([{}] * (s - len(pred_probs[b]) + 1))
+
+            pred_probs[b] = pred_probs[b][:seg_len]
+
+        pred_probs = [[p for p in batch if p != {}] for batch in pred_probs]
+        
+        return pred_chars, pred_probs
+
+    def post_normalize(self, text):
+        text = re.sub(r'\s*[。\.]', '.', text)
+        text = re.sub(r'\s*[，、,]', ',', text)
+        text = re.sub(r'\s*[！!]', '!', text)
+        text = re.sub(r'\s*[？\?]', '?', text)
+        return text

special_tokens_map.json

@@ -0,0 +1 @@
+{}

tokenization_nombert.py

@@ -0,0 +1,64 @@
+import json
+import os
+from transformers import PreTrainedTokenizer
+
+
+class NomTokenizer(PreTrainedTokenizer):
+    vocab_files_names = {'vocab_file': 'vocab.json'}
+    
+    def __init__(
+        self,
+        vocab_file,
+        unk_token='<UNK>',
+        unk_token_id=0,
+        id_start=1,
+        **kwargs
+    ):
+        self.vocab_file = vocab_file
+        self.id_start = id_start
+        self.unk_token = unk_token
+        self.unk_token_id = unk_token_id
+        self.pad_token = unk_token
+        self.pad_token_id = unk_token_id
+
+        with open(vocab_file, 'r', encoding='utf-8') as f:
+            self.vocab_dict = json.load(f)
+    
+        self.char2id = {}
+        self.id2char = {}
+        for i, char in enumerate(self.vocab_dict.keys(), start=id_start):
+            self.char2id[char] = i
+            self.id2char[i] = char
+        self.id_to_options = {idx: v for idx, v in enumerate(self.vocab_dict.values(), start=id_start)}
+
+        super().__init__(**kwargs)
+
+    def _tokenize(self, text):
+        return list(text)
+
+    def _convert_token_to_id(self, token):
+        return self.char2id.get(token, self.unk_token_id)
+
+    def _convert_id_to_token(self, index):
+        if index == self.unk_token_id:
+            return self.unk_token
+        return self.id2char.get(index, self.unk_token)
+
+    @property
+    def vocab_size(self):
+        return len(self.char2id) + 1
+
+    def get_vocab(self):
+        vocab = {**self.char2id, **self.added_tokens_encoder}
+        return vocab
+
+    def save_vocabulary(self, save_directory, filename_prefix=None):
+        if filename_prefix:
+            vocab_file = os.path.join(save_directory, f'{filename_prefix}-vocab.json')
+        else:
+            vocab_file = os.path.join(save_directory, 'vocab.json')
+        
+        with open(vocab_file, 'w', encoding='utf-8') as f:
+            json.dump(self.vocab_dict, f, ensure_ascii=False)
+        
+        return (vocab_file,)

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "added_tokens_decoder": {},
+  "clean_up_tokenization_spaces": false,
+  "extra_special_tokens": {},
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "NomTokenizer",
+  "auto_map": {
+    "AutoTokenizer": [
+      "tokenization_nombert.NomTokenizer",
+      null
+      ]
+  },
+  "unk_token": "<UNK>"
+}