class_docstr: BERT Regression model.
class_docstr_tok: ["BERT","Regression","model","."]
code: def build_flax_model(self):
      model_dtype = getattr(
          jnp, self.config.get("model_dtype_str", "float32")
      )
      with self.config.unlocked():
          # Add vocabulary information from dataset meta-data to configs:
          self.config.model.stem.vocab_size = self.dataset_meta_data[
              "vocab_size"
          ]
          self.config.model.stem.type_vocab_size = (
              self.dataset_meta_data["type_vocab_size"]
          )
      return BERT(
          stem_config=self.config.model.stem,
          encoder_config=self.config.model.encoder,
          head_config=self.config.model.head,
          dtype=model_dtype,
      )
  
imports: from typing import Any, Dict
  from absl import logging
  import flax
  import flax.linen
  import jax.numpy
  import ml_collections
  from scenic.common_lib import debug_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.base_models.regression_model import RegressionModel
  from scenic.projects.baselines.bert import bert_base_model
  from scenic.projects.baselines.bert import layers
label_desc: Use image node_id 1 for calling the BERTRegressionModel obj's underlying member method code with example usage: obj.build_flax_model() and returns: BERT