code: def model():
      """
      Rebuild XLMR model from training script but with reduces layer count to speed up unit test
      """
      num_classes = 2
      input_dim = 768
  
      # Would be more elegant to mock RobertaEncoderConf and default num_encoder_layers to 1 but I failed do far
      from torchtext.models import RobertaClassificationHead
      from torchtext.models.roberta.bundler import (
          _TEXT_BUCKET,
          RobertaEncoderConf,
          T,
          load_state_dict_from_url,
          urljoin,
      )
  
      try:
          import torchtext.models.roberta.bundler.RobertaModelBundle as RobertaBundle
      except ImportError:
          from torchtext.models.roberta.bundler import RobertaBundle
  
      torch.manual_seed(42)
  
      XLMR_BASE_ENCODER = RobertaBundle(
          _path=urljoin(_TEXT_BUCKET, "xlmr.base.encoder.pt"),
          _encoder_conf=RobertaEncoderConf(
              vocab_size=250002, num_encoder_layers=1
          ),
          transform=lambda: T.Sequential(
              T.SentencePieceTokenizer(
                  urljoin(_TEXT_BUCKET, "xlmr.sentencepiece.bpe.model")
              ),
              T.VocabTransform(
                  load_state_dict_from_url(
                      urljoin(_TEXT_BUCKET, "xlmr.vocab.pt")
                  )
              ),
              T.Truncate(254),
              T.AddToken(token=0, begin=True),
              T.AddToken(token=2, begin=False),
          ),
      )
  
      classifier_head = RobertaClassificationHead(
          num_classes=num_classes, input_dim=input_dim
      )
      model = XLMR_BASE_ENCODER.get_model(
          head=classifier_head, load_weights=False
      )
  
      yield model
  
imports: import importlib
  import json
  import os
  import shutil
  import sys
  import time
  from argparse import Namespace
  from pathlib import Path
  import pytest
  import requests
  import test_utils
  import torch
  from model_archiver import ModelArchiverConfig
  from test_utils import REPO_ROOT
  from ts.torch_handler.unit_tests.test_utils.mock_context import MockContext
label_desc: Use image node_id 2 for calling a global function with example usage: model() without return types