Add MobileBERT testing on riscv64 (fp32 + quantize)

Author: luhenry

.github/workflows/riscv64.yml

@@ -28,12 +28,14 @@ jobs:
     strategy:
       matrix:
         include:
-          - { model: add, xnnpack: false, quantize: false }
-          - { model: add, xnnpack: true,  quantize: false }
-          - { model: mv2, xnnpack: false, quantize: false }
-          - { model: mv2, xnnpack: false, quantize: true }
-          - { model: mv2, xnnpack: true,  quantize: false }
-          - { model: mv2, xnnpack: true,  quantize: true }
+          - { model: add,        xnnpack: false, quantize: false }
+          - { model: add,        xnnpack: true,  quantize: false }
+          - { model: mv2,        xnnpack: false, quantize: false }
+          - { model: mv2,        xnnpack: true,  quantize: false }
+          - { model: mv2,        xnnpack: true,  quantize: true }
+          - { model: mobilebert, xnnpack: false, quantize: false }
+          - { model: mobilebert, xnnpack: true,  quantize: false }
+          - { model: mobilebert, xnnpack: true,  quantize: true }
     permissions:
       id-token: write
       contents: read

examples/riscv/aot_riscv.py

@@ -49,7 +49,38 @@ def build_mv2():
     return model, example_inputs, test_inputs, False
 
 
-MODELS = {"add": build_add, "mv2": build_mv2}
+def build_mobilebert():
+    from transformers import MobileBertConfig, MobileBertModel
+
+    config = MobileBertConfig(
+        vocab_size=1024,
+        hidden_size=128,
+        embedding_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=2,
+        intermediate_size=128,
+        intra_bottleneck_size=32,
+    )
+
+    class Wrapper(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.model = MobileBertModel(config).eval()
+
+        def forward(self, input_ids):
+            return self.model(input_ids).last_hidden_state
+
+    model = Wrapper().eval()
+    example_inputs = (torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8]]),)
+    test_inputs = [example_inputs]
+    return model, example_inputs, test_inputs, False
+
+
+MODELS = {
+    "add": build_add,
+    "mv2": build_mv2,
+    "mobilebert": build_mobilebert,
+}
 
 
 def main() -> None:

examples/riscv/requirements.txt

@@ -1 +1,2 @@
 torchvision
+transformers