class_docstr: Implements a video+text single-tower backbone.
  
  Attributes:
    num_classes: Number of output classes.
    video_tower_config: The config of the video tower.
    text_tower_config: The config of the text tower.
    video_text_fusion_config: The config of video+text fusion.
    classifier: 'gap' or 'token'
class_docstr_tok: ["Implements","a","video+text","single-tower","backbone",".","Attributes",":","num_classes",":","Number","of","output","classes",".","video_tower_config",":","The","config","of","the","video","tower",".","text_tower_config",":","The","config","of","the","text","tower",".","video_text_fusion_config",":","The","config","of","video+text","fusion",".","classifier",":","'gap","'","or","'token","'"]
docstr: Runs model inference.
  
  In this model, the video encoder, text encoder, and video-text fusion
  encoder are shared among all tasks. We build a different head for a
  different task and/or a different dataset. The head will have a unique name
  '{dataset}_{task}_head' if dataset is given. Otherwise, the head will have
  a name '{task}_head'.
  
  Args:
    inputs: Input dict containing the rgb frames or rgb embeddings and
      tokenized texts or text embeddings. RGB frames has a shape (batch_size,
      num_frames, height, width, channels) and RGB embeddings has a shape
      (batch_size, num_frames, channels). inputs['class_names'] or
      inputs['caption'] is dict of three elements whose keys are 'input_mask',
      'input_word_ids', and 'input_type_ids'.
    task: 'action_segmentation', classification', 'temporal_localization',
      'moment_retrieval' or 'highlight_detection'.
    dataset: The name of the dataset. The name will be used to create
      different heads for different datasets.
    train: Whether or not the model is under training.
    debug: Whether or not it is in debug mode.
  
  Returns:
    A 2D float tensor of shape (batch_size, num_classes) representing the
    logits for each class.
docstr_tok: ["Runs","model","inference",".","In","this","model",",","the","video","encoder",",","text","encoder",",","and","video-text","fusion","encoder","are","shared","among","all","tasks",".","We","build","a","different","head","for","a","different","task","and\/or","a","different","dataset",".","The","head","will","have","a","unique","name","'","{","dataset","}","_","{","task","}","_head","'","if","dataset","is","given",".","Otherwise",",","the","head","will","have","a","name","'","{","task","}","_head","'",".","Args",":","inputs",":","Input","dict","containing","the","rgb","frames","or","rgb","embeddings","and","tokenized","texts","or","text","embeddings",".","RGB","frames","has","a","shape","(","batch_size",",","num_frames",",","height",",","width",",","channels",")","and","RGB","embeddings","has","a","shape","(","batch_size",",","num_frames",",","channels",")",".","inputs","[","'class_names","'","]","or","inputs","[","'caption","'","]","is","dict","of","three","elements","whose","keys","are","'input_mask","'",",","'input_word_ids","'",",","and","'input_type_ids","'",".","task",":","'action_segmentation","'",",","classification","'",",","'temporal_localization","'",",","'moment_retrieval","'","or","'highlight_detection","'",".","dataset",":","The","name","of","the","dataset",".","The","name","will","be","used","to","create","different","heads","for","different","datasets",".","train",":","Whether","or","not","the","model","is","under","training",".","debug",":","Whether","or","not","it","is","in","debug","mode",".","Returns",":","A","2D","float","tensor","of","shape","(","batch_size",",","num_classes",")","representing","the","logits","for","each","class","."]
code: def __call__(
      self,
      inputs: Dict[str, Any],
      task: str = "classification",
      dataset: str = "",
      train: bool = False,
      debug: bool = False,
  ):
      """Runs model inference.
  
      In this model, the video encoder, text encoder, and video-text fusion
      encoder are shared among all tasks. We build a different head for a
      different task and/or a different dataset. The head will have a unique name
      '{dataset}_{task}_head' if dataset is given. Otherwise, the head will have
      a name '{task}_head'.
  
      Args:
        inputs: Input dict containing the rgb frames or rgb embeddings and
          tokenized texts or text embeddings. RGB frames has a shape (batch_size,
          num_frames, height, width, channels) and RGB embeddings has a shape
          (batch_size, num_frames, channels). inputs['class_names'] or
          inputs['caption'] is dict of three elements whose keys are 'input_mask',
          'input_word_ids', and 'input_type_ids'.
        task: 'action_segmentation', classification', 'temporal_localization',
          'moment_retrieval' or 'highlight_detection'.
        dataset: The name of the dataset. The name will be used to create
          different heads for different datasets.
        train: Whether or not the model is under training.
        debug: Whether or not it is in debug mode.
  
      Returns:
        A 2D float tensor of shape (batch_size, num_classes) representing the
        logits for each class.
      """
      assert task in {
          "action_segmentation",
          "classification",
          "temporal_localization",
          "moment_retrieval",
          "highlight_detection",
      }
  
      input_word_ids = None
      text_input_mask = None
      if self.text_tower_config is not None:
          if task == "moment_retrieval":
              # Merges all captions
              text_inputs = jax.tree_util.tree_map(
                  lambda x: x.reshape((-1, x.shape[-1])),
                  inputs["caption"],
              )
          elif task == "highlight_detection":
              input_key = self.text_tower_config.get(
                  "input_key", "video_title"
              )
              text_inputs = inputs[input_key]
          else:
              # The class names are the same in all batches.
              text_inputs = jax.tree_util.tree_map(
                  lambda x: x[0], inputs["class_names"]
              )
          if (
              self.text_tower_config.get("input_type", "tokenized_text")
              == "tokenized_text"
          ):
              if task in {
                  "action_segmentation",
                  "temporal_localization",
                  "classification",
              }:
                  assert (
                      text_inputs["input_word_ids"].shape[0]
                      == self.num_classes
                  )
              input_word_ids = text_inputs["input_word_ids"]
              text_input_mask = text_inputs["input_mask"]
      video_tokens = self.encode_video(inputs, train=train, debug=debug)
      text_tokens = self.encode_text(
          inputs, task=task, train=train, debug=debug
      )
      return self.fuse_video_text(
          video_tokens,
          text_tokens,
          task,
          input_word_ids,
          text_input_mask,
          video_input_mask=inputs.get("input_mask"),
          train=train,
          debug=debug,
      )
  
imports: from typing import Any, Dict, Optional
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.projects.unloc import action_segmentation_base_model
  from scenic.projects.unloc import encoders
  from scenic.projects.unloc import heads
  from scenic.projects.unloc import moment_retrieval_base_model
  from scenic.projects.unloc import temporal_localization_base_model
  from scenic.projects.unloc import video_text_fusion
label_desc: Use image node_id 5 for calling the VideoTextSingleTower obj's underlying member method code with example usage: obj.__call__(inputs, task, dataset, train, debug) and returns: self