[frontend] add tiny bert

examples/chinese_prosody_polyphone/README.md

@@ -31,12 +31,16 @@ TODO(Binbin Zhang): Add more data
 
 BERT-MLT is for polyphone and prosody joint training.
 
+Bert base model: https://huggingface.co/google-bert/bert-base-chinese
+TinyBert base model: https://huggingface.co/JeremiahZ/TinyBERT_4L_zh_backup/tree/main
+
 ### Polyphone
 
 | system         | ACC    |
 |----------------|--------|
 | BERT-polyphone | 0.9778 |
 | BERT-MLT       | 0.9797 |
+| TinyBert-4L-MLT| 0.9736 |
 
 
 ### Prosody
@@ -45,5 +49,6 @@ BERT-MLT is for polyphone and prosody joint training.
 |---------------------------|--------|--------|--------|
 | BERT-prosody              | 0.9308 | 0.8058 | 0.8596 |
 | BERT-MLT                  | 0.9334 | 0.8088 | 0.8559 |
+| TinyBert-4L-MLT           | 0.9241 | 0.7769 | 0.7829 |
 | BERT-prosody (exclude #4) | 0.9233 | 0.7074 | 0.6120 |
 | BERT-MLT (exclude #4)     | 0.9261 | 0.7146 | 0.6140 |

examples/chinese_prosody_polyphone/run.sh

@@ -5,7 +5,8 @@ stage=0
 stop_stage=4
 url=https://wetts-1256283475.cos.ap-shanghai.myqcloud.com/data
 
-dir=exp
+bert_name_or_path=bert-chinese-base
+dir=exp/bert
 
 . tools/parse_options.sh
 
@@ -35,6 +36,7 @@ fi
 if [ ${stage} -le 1 ] && [ ${stop_stage} -ge 1 ]; then
   mkdir -p $dir
   python frontend/train.py \
+    --bert_name_or_path $bert_name_or_path \
     --gpu 2 \
     --lr 0.001 \
     --num_epochs 10 \
@@ -54,6 +56,7 @@ fi
 if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
   # Test polyphone, metric: accuracy
   python frontend/test_polyphone.py \
+    --bert_name_or_path $bert_name_or_path \
     --polyphone_dict lexicon/polyphone.txt \
     --prosody_dict lexicon/prosody.txt \
     --test_data data/polyphone/test.txt \
@@ -62,6 +65,7 @@ if [ ${stage} -le 2 ] && [ ${stop_stage} -ge 2 ]; then
 
   # Test prosody, metric: F1-score
   python frontend/test_prosody.py \
+    --bert_name_or_path $bert_name_or_path \
     --polyphone_dict lexicon/polyphone.txt \
     --prosody_dict lexicon/prosody.txt \
     --test_data data/prosody/cv.txt \
@@ -73,6 +77,7 @@ fi
 if [ ${stage} -le 3 ] && [ ${stop_stage} -ge 3 ]; then
   # export onnx model
   python frontend/export_onnx.py \
+    --bert_name_or_path $bert_name_or_path \
     --polyphone_dict lexicon/polyphone.txt \
     --prosody_dict lexicon/prosody.txt \
     --checkpoint $dir/9.pt \
@@ -85,8 +90,9 @@ if [ ${stage} -le 4 ] && [ ${stop_stage} -ge 4 ]; then
   # pinyin ['ba1', 'fang1', 'cai2', 'bao3', 'jin4']
   # prosody [0 1 0 0 4]
   python frontend/g2p_prosody.py \
+    --bert_name_or_path $bert_name_or_path \
     --text "八方财宝进" \
-    --hanzi2pinyin_file lexicon/pinyin_dict.txt \
+    --lexicon_file lexicon/pinyin_dict.txt \
     --polyphone_file lexicon/polyphone.txt \
     --polyphone_prosody_model $dir/9.onnx
 fi

wetts/frontend/dataset.py

@@ -15,21 +15,21 @@
 import numpy as np
 import torch
 from torch.utils.data import Dataset
-from transformers import AutoTokenizer
 
 IGNORE_ID = -100
 
-tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")
-
 
 class FrontendDataset(Dataset):
+
     def __init__(
         self,
+        tokenizer=None,
         polyphone_file=None,
         polyphone_dict=None,
         prosody_file=None,
         prosody_dict=None,
     ):
+        self.tokenizer = tokenizer
         self.data = []
         self.polyphone_dict = polyphone_dict
         self.prosody_dict = prosody_dict
@@ -58,22 +58,21 @@ def read_polyphone_data(self, polyphone_file):
                 tokens = []
                 all_polyphones = []
                 for i in range(0, len(arr), 2):
-                    toks = tokenizer.encode(arr[i], add_special_tokens=False)
+                    toks = self.tokenizer.encode(arr[i],
+                                                 add_special_tokens=False)
                     polyphones = [IGNORE_ID] * len(toks)
                     if i + 1 < len(arr):
                         if arr[i + 1] in self.polyphone_dict:
                             polyphones[-1] = self.polyphone_dict[arr[i + 1]]
-                        else:
-                            print("Skip unknown polyphone {}".format(arr[i + 1]))
+                        # else:
+                        #     print("Skip unknown polyphone {}".format(arr[i + 1]))
                     all_polyphones.extend(polyphones)
                     tokens.append(arr[i])
-                data.append(
-                    {
-                        "sentence": tokens,
-                        "polyphones": all_polyphones,
-                        "prosody": [IGNORE_ID] * len(all_polyphones),
-                    }
-                )
+                data.append({
+                    "sentence": tokens,
+                    "polyphones": all_polyphones,
+                    "prosody": [IGNORE_ID] * len(all_polyphones),
+                })
         return data
 
     def read_prosody_data(self, prosody_file):
@@ -93,11 +92,10 @@ def read_prosody_data(self, prosody_file):
                 else:
                     for i in range(0, len(arr), 2):
                         # #N and N is less than `prosody_dict`
-                        if (
-                            arr[i + 1][0] != "#"
-                            or not arr[i + 1][1:].encode("utf8").isdigit()
-                            or int(arr[i + 1][1:]) >= len(self.prosody_dict)
-                        ):
+                        if (arr[i + 1][0] != "#"
+                                or not arr[i + 1][1:].encode("utf8").isdigit()
+                                or int(arr[i + 1][1:]) >= len(
+                                    self.prosody_dict)):
                             ok = False
                             break
                 if not ok:
@@ -107,19 +105,18 @@ def read_prosody_data(self, prosody_file):
                 tokens = []
                 prosody = []
                 for i in range(0, len(arr), 2):
-                    toks = tokenizer.encode(arr[i], add_special_tokens=False)
+                    toks = self.tokenizer.encode(arr[i],
+                                                 add_special_tokens=False)
                     rank = int(arr[i + 1][1:])
                     rhythm = [0] * len(toks)
                     rhythm[-1] = rank
                     prosody.extend(rhythm)
                     tokens.append(arr[i])
-                data.append(
-                    {
-                        "sentence": tokens,
-                        "polyphones": [IGNORE_ID] * len(prosody),
-                        "prosody": prosody,
-                    }
-                )
+                data.append({
+                    "sentence": tokens,
+                    "polyphones": [IGNORE_ID] * len(prosody),
+                    "prosody": prosody,
+                })
         return data
 
     def __len__(self):
@@ -129,7 +126,7 @@ def __getitem__(self, idx):
         return self.data[idx]
 
 
-def collote_fn(batch_samples):
+def collate_fn(batch_samples, tokenizer):
     batch_sentence, batch_polyphones, batch_prosody = [], [], []
     for sample in batch_samples:
         batch_sentence.append(sample["sentence"])
@@ -142,13 +139,16 @@ def collote_fn(batch_samples):
         is_split_into_words=True,
         return_tensors="pt",
     )
-    polyphone_label = np.ones_like(batch_inputs["input_ids"], dtype=int) * IGNORE_ID
+    polyphone_label = np.ones_like(batch_inputs["input_ids"],
+                                   dtype=int) * IGNORE_ID
     for idx, polyphones in enumerate(batch_polyphones):
-        polyphone_label[idx][1 : len(polyphones) + 1] = np.array(
-            polyphones, dtype=np.int32
-        )
-    prosody_label = np.ones_like(batch_inputs["input_ids"], dtype=int) * IGNORE_ID
+        polyphone_label[idx][1:len(polyphones) + 1] = np.array(polyphones,
+                                                               dtype=np.int32)
+    prosody_label = np.ones_like(batch_inputs["input_ids"],
+                                 dtype=int) * IGNORE_ID
     for idx, prosody in enumerate(batch_prosody):
-        prosody_label[idx][1 : len(prosody) + 1] = np.array(prosody, dtype=np.int32)
+        prosody_label[idx][1:len(prosody) + 1] = np.array(prosody,
+                                                          dtype=np.int32)
 
-    return batch_inputs, torch.tensor(polyphone_label), torch.tensor(prosody_label)
+    return batch_inputs, torch.tensor(polyphone_label), torch.tensor(
+        prosody_label)

wetts/frontend/export_onnx.py

@@ -14,20 +14,40 @@
 
 import argparse
 import os
+import random
 
-import torch
 import onnxruntime as ort
-
+import torch
+import torch.nn.functional as F
 from model import FrontendModel
+from onnxruntime.quantization import QuantType, quantize_dynamic
+from transformers import AutoTokenizer
 from utils import read_table
 
 
+def cosine_similarity(tensor1, tensor2):
+    flat1 = tensor1.flatten()
+    flat2 = tensor2.flatten()
+    similarity = F.cosine_similarity(flat1.unsqueeze(0),
+                                     flat2.unsqueeze(0),
+                                     dim=1)
+    return similarity.item()
+
+
 def get_args():
     parser = argparse.ArgumentParser(description="export onnn model")
-    parser.add_argument("--polyphone_dict", required=True, help="polyphone dict file")
-    parser.add_argument("--prosody_dict", required=True, help="train data file")
+    parser.add_argument("--polyphone_dict",
+                        required=True,
+                        help="polyphone dict file")
+    parser.add_argument("--prosody_dict",
+                        required=True,
+                        help="train data file")
     parser.add_argument("--checkpoint", required=True, help="checkpoint model")
     parser.add_argument("--onnx_model", required=True, help="onnx model path")
+    parser.add_argument("--quant_onnx_model", help="quant onnx model path")
+    parser.add_argument("--bert_name_or_path",
+                        default='bert-chinese-base',
+                        help="bert init model")
     args = parser.parse_args()
     return args
 
@@ -41,8 +61,10 @@ def main():
     num_prosody = len(prosody_dict)
 
     # Init model
-    model = FrontendModel(num_polyphones, num_prosody)
-    model.load_state_dict(torch.load(args.checkpoint, map_location="cpu"))
+    tokenizer = AutoTokenizer.from_pretrained(args.bert_name_or_path)
+    model = FrontendModel(num_polyphones, num_prosody, args.bert_name_or_path)
+    model.load_state_dict(torch.load(args.checkpoint, map_location="cpu"),
+                          strict=False)
     model.forward = model.export_forward
     model.eval()
 
@@ -54,29 +76,40 @@ def main():
         input_names=["input"],
         output_names=["polyphone_output", "prosody_output"],
         dynamic_axes={
-            "input": {1: "T"},
-            "polyphone_output": {1: "T"},
-            "prosody_output": {1: "T"},
+            "input": {
+                1: "T"
+            },
+            "polyphone_output": {
+                1: "T"
+            },
+            "prosody_output": {
+                1: "T"
+            },
         },
         opset_version=13,
         verbose=False,
     )
 
     # Verify onnx precision
-    torch_output = model(dummy_input)
-    ort_sess = ort.InferenceSession(args.onnx_model)
-    onnx_output = ort_sess.run(None, {"input": dummy_input.numpy()})
-    print(torch_output[1])
-    print(onnx_output[1])
-    if torch.allclose(
-        torch_output[0], torch.tensor(onnx_output[0]), atol=1e-3
-    ) and torch.allclose(torch_output[1], torch.tensor(onnx_output[1]), atol=1e-3):
-        print("Export to onnx succeed!")
-    else:
-        print(
-            """Export to onnx succeed, but pytorch/onnx have different
-                 outputs when given the same input, please check!!!"""
-        )
+    rand_len = random.randint(5, 20)
+    rand_input = torch.randint(0, 100, (1, rand_len))
+    print('rand input', rand_input)
+    torch_output = model(rand_input)
+    print(torch_output[0].size(), torch_output[1].size())
+
+    def verify_export(onnx_model):
+        ort_sess = ort.InferenceSession(onnx_model)
+        onnx_output = ort_sess.run(None, {"input": rand_input.numpy()})
+        sim = cosine_similarity(torch_output[1],
+                                torch.from_numpy(onnx_output[1]))
+        print(f"Export to onnx {onnx_model}, similarity {sim}")
+
+    verify_export(args.onnx_model)
+    if args.quant_onnx_model is not None:
+        quantize_dynamic(args.onnx_model,
+                         args.quant_onnx_model,
+                         weight_type=QuantType.QUInt8)
+        verify_export(args.quant_onnx_model)
 
 
 if __name__ == "__main__":