xref: /aosp_15_r20/external/executorch/examples/models/mobilebert/model.py (revision 523fa7a60841cd1ecfb9cc4201f1ca8b03ed023a) 
1# Copyright (c) Meta Platforms, Inc. and affiliates.
2# All rights reserved.
3#
4# This source code is licensed under the BSD-style license found in the
5# LICENSE file in the root directory of this source tree.
6
7import logging
8
9import torch
10
11from transformers import AutoTokenizer, MobileBertModel  # @manual
12
13from ..model_base import EagerModelBase
14
15
16class MobileBertModelExample(EagerModelBase):
17    def __init__(self):
18        pass
19
20    def get_eager_model(self) -> torch.nn.Module:
21        logging.info("loading mobilebert model")
22        # pyre-ignore
23        model = MobileBertModel.from_pretrained(
24            "google/mobilebert-uncased", return_dict=False
25        )
26        logging.info("loaded mobilebert model")
27        return model
28
29    def get_example_inputs(self):
30        tokenizer = AutoTokenizer.from_pretrained("google/mobilebert-uncased")
31        return (tokenizer("Hello, my dog is cute", return_tensors="pt")["input_ids"],)
32