Support bert distilbert (#14138)

### Summary
Add export scripts for supporting bert distilbert

---------

Co-authored-by: neuropilot-captain <[email protected]>

Author: neuropilot-captain

examples/mediatek/README.md

@@ -71,7 +71,7 @@ source shell_scripts/export_llama.sh <model_name> <num_chunks> <prompt_num_token
 bash shell_scripts/export_oss.sh <model_name>
 ```
 - Argument Options:
-    - `model_name`: deeplabv3/edsr/inceptionv3/inceptionv4/mobilenetv2/mobilenetv3/resnet18/resnet50
+    - `model_name`: deeplabv3/edsr/inceptionv3/inceptionv4/mobilenetv2/mobilenetv3/resnet18/resnet50/dcgan/wav2letter/vit_b_16/mobilebert/emformer_rnnt/bert/distilbert
 
 # Runtime
 ## Environment Setup

examples/mediatek/aot_utils/oss_utils/utils.py

@@ -5,9 +5,11 @@
 # LICENSE file in the root directory of this source tree.
 
 import os
+import random
 from typing import Optional
 
 import torch
+import transformers
 from executorch import exir
 from executorch.backends.mediatek import (
     NeuropilotPartitioner,
@@ -42,6 +44,7 @@ def build_executorch_binary(
         quantized_model = convert_pt2e(annotated_model, fold_quantize=False)
         aten_dialect = torch.export.export(quantized_model, inputs, strict=True)
     else:
+        print("Using float model...")
         aten_dialect = torch.export.export(model, inputs, strict=True)
 
     from executorch.exir.program._program import to_edge_transform_and_lower
@@ -71,3 +74,58 @@ def make_output_dir(path: str):
             os.remove(os.path.join(path, f))
         os.removedirs(path)
     os.makedirs(path)
+
+
+def get_masked_language_model_dataset(dataset_path, tokenizer, data_size, shuffle=True):
+
+    def get_data_loader():
+        class MaskedSentencesDataset(torch.utils.data.Dataset):
+            def __init__(self, dataset_path, tokenizer, data_size) -> None:
+                self.data_size = data_size
+                self.dataset = self._get_val_dataset(dataset_path, data_size, tokenizer)
+
+            def _get_val_dataset(self, dataset_path, data_size, tokenizer):
+                data_collator = transformers.DataCollatorForLanguageModeling(
+                    tokenizer=tokenizer
+                )
+                with open(dataset_path, "r") as f:
+                    texts = f.read().split("\n")
+                    texts = [
+                        text for text in random.choices(texts, k=2000) if len(text) > 1
+                    ]
+                    dataset = data_collator([tokenizer(text) for text in texts])
+                return dataset
+
+            def __getitem__(self, idx):
+                return (
+                    self.dataset["input_ids"][idx].to(torch.int32),
+                    self.dataset["attention_mask"][idx].to(torch.float32),
+                    self.dataset["labels"][idx],
+                )
+
+            def __len__(self):
+                return self.data_size
+
+        dataset = MaskedSentencesDataset(dataset_path, tokenizer, data_size)
+        return torch.utils.data.DataLoader(
+            dataset,
+            shuffle=shuffle,
+        )
+
+    # prepare input data
+    inputs, targets = [], []
+    data_loader = get_data_loader()
+    for data in data_loader:
+        if len(inputs) >= data_size:
+            break
+        input_ids = data[0]
+        attention_mask = data[1]
+        target = data[2][0]
+        indice = [i for i, x in enumerate(target) if x != -100]
+        # continue if no mask annotated
+        if len(indice) == 0:
+            continue
+        inputs.append((input_ids, attention_mask))
+        targets.append(target)
+
+    return inputs, targets

examples/mediatek/model_export_scripts/bert.py

@@ -0,0 +1,82 @@
+# Copyright (c) MediaTek Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import os
+import sys
+
+if os.getcwd() not in sys.path:
+    sys.path.append(os.getcwd())
+
+from aot_utils.oss_utils.utils import (
+    build_executorch_binary,
+    get_masked_language_model_dataset,
+)
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+
+def main(args):
+    # ensure the working directory exist.
+    os.makedirs(args.artifact, exist_ok=True)
+    data_size = 100
+
+    tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
+    inputs, targets = get_masked_language_model_dataset(
+        args.dataset, tokenizer, data_size
+    )
+
+    # build pte
+    module = AutoModelForMaskedLM.from_pretrained(
+        "google-bert/bert-base-uncased"
+    ).eval()
+    pte_filename = "bert_mtk"
+
+    build_executorch_binary(
+        module,
+        inputs[0],
+        f"{args.artifact}/{pte_filename}",
+        inputs,
+        skip_op_name={"aten_embedding_default", "aten_where_self"},
+    )
+
+    # save data to inference on device
+    input_list_file = f"{args.artifact}/input_list.txt"
+    with open(input_list_file, "w") as f:
+        for i in range(len(inputs)):
+            f.write(f"input_{i}_0.bin input_{i}_1.bin\n")
+    for idx, data in enumerate(inputs):
+        for i, d in enumerate(data):
+            file_name = f"{args.artifact}/input_{idx}_{i}.bin"
+            d.detach().numpy().tofile(file_name)
+    for idx, data in enumerate(targets):
+        file_name = f"{args.artifact}/golden_{idx}_0.bin"
+        data.detach().numpy().tofile(file_name)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-a",
+        "--artifact",
+        help="path for storing generated artifacts and output by this example. Default ./bert",
+        default="./bert",
+        type=str,
+    )
+    parser.add_argument(
+        "-d",
+        "--dataset",
+        help=(
+            "path to the validation text. "
+            "e.g. --dataset wikisent2.txt "
+            "for https://www.kaggle.com/datasets/mikeortman/wikipedia-sentences"
+        ),
+        default="wikisent2.txt",
+        type=str,
+        required=False,
+    )
+
+    args = parser.parse_args()
+    main(args)

examples/mediatek/model_export_scripts/distilbert.py

@@ -0,0 +1,82 @@
+# Copyright (c) MediaTek Inc.
+# All rights reserved
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import os
+import sys
+
+if os.getcwd() not in sys.path:
+    sys.path.append(os.getcwd())
+
+from aot_utils.oss_utils.utils import (
+    build_executorch_binary,
+    get_masked_language_model_dataset,
+)
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+
+def main(args):
+    # ensure the working directory exist.
+    os.makedirs(args.artifact, exist_ok=True)
+    data_size = 100
+
+    tokenizer = AutoTokenizer.from_pretrained("distilbert/distilbert-base-uncased")
+    inputs, targets = get_masked_language_model_dataset(
+        args.dataset, tokenizer, data_size
+    )
+
+    # build pte
+    module = AutoModelForMaskedLM.from_pretrained(
+        "distilbert/distilbert-base-uncased"
+    ).eval()
+    pte_filename = "distilbert_mtk"
+
+    build_executorch_binary(
+        module,
+        inputs[0],
+        f"{args.artifact}/{pte_filename}",
+        inputs,
+        skip_op_name={"aten_embedding_default", "aten_where_self"},
+    )
+
+    # save data to inference on device
+    input_list_file = f"{args.artifact}/input_list.txt"
+    with open(input_list_file, "w") as f:
+        for i in range(len(inputs)):
+            f.write(f"input_{i}_0.bin input_{i}_1.bin\n")
+    for idx, data in enumerate(inputs):
+        for i, d in enumerate(data):
+            file_name = f"{args.artifact}/input_{idx}_{i}.bin"
+            d.detach().numpy().tofile(file_name)
+    for idx, data in enumerate(targets):
+        file_name = f"{args.artifact}/golden_{idx}_0.bin"
+        data.detach().numpy().tofile(file_name)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-a",
+        "--artifact",
+        help="path for storing generated artifacts and output by this example. Default ./distilbert",
+        default="./distilbert",
+        type=str,
+    )
+    parser.add_argument(
+        "-d",
+        "--dataset",
+        help=(
+            "path to the validation text. "
+            "e.g. --dataset wikisent2.txt "
+            "for https://www.kaggle.com/datasets/mikeortman/wikipedia-sentences"
+        ),
+        default="wikisent2.txt",
+        type=str,
+        required=False,
+    )
+
+    args = parser.parse_args()
+    main(args)

examples/mediatek/shell_scripts/export_oss.sh

@@ -41,4 +41,10 @@ then
 elif [ $model = "emformer_rnnt" ]
 then
 	python3 model_export_scripts/emformer_rnnt.py
+elif [ $model = "bert" ]
+then
+	python3 model_export_scripts/bert.py
+elif [ $model = "distilbert" ]
+then
+	python3 model_export_scripts/distilbert.py
 fi