class_docstr: Control the loops for training, evaluation, and prediction. This module is independent of
  the model definition. This class inherits from the Pytorch Lightning's LightningModule:
  https://lightning.ai/docs/pytorch/stable/common/lightning_module.html
class_docstr_tok: ["Control","the","loops","for","training",",","evaluation",",","and","prediction",".","This","module","is","independent","of","the","model","definition",".","This","class","inherits","from","the","Pytorch","Lightning","'s","LightningModule",":","https",":","\/\/lightning.ai\/docs\/pytorch\/stable\/common\/lightning_module.html"]
docstr: Parameters
  ----------
  query_model
      The query model.
  response_model
      The response model.
  signature
      query or response.
  match_label
      The label of match class.
  matches
      A list of DictConfigs, each of which defines one pair of feature column match and the configs
      to compute the matching loss.
  optim_type
      Optimizer type. We now support:
      - adamw
      - adam
      - sgd
  lr_choice
      How to set each layer's learning rate. If not specified, the default is a single
      learnng rate for all layers. Otherwise, we now support two choices:
      - two_stages
          The layers in the pretrained models have a small learning rate (lr * lr_mult),
          while the newly added head layers use the provided learning rate.
      - layerwise_decay
          The layers have decreasing learning rate from the output end to the input end.
          The intuition is that later layers are more task-related, hence larger learning rates.
  lr_schedule
      Learning rate schedule. We now support:
      - cosine_decay
          Linear warmup followed by cosine decay
      - polynomial_decay
          Linear warmup followed by polynomial decay
  lr
      Learning rate.
  lr_decay
      The learning rate decay factor (0, 1). It is used only when lr_choice is "layerwise_decay".
  end_lr
      The final learning rate after decay.
  lr_mult
      The learning rate multiplier (0, 1). It is used only when lr_choice is "two_stages".
  weight_decay
      The weight decay to regularize layer weights' l2 norm.
  warmup_steps
      How many steps to warmup learning rate. If a float (0, 1), it would represent the
      percentage of steps over all the training steps. The actual number is calculated as
      "int(warmup_steps * max_steps)". If an integer, it would be the exact step number.
  loss_func
      A Pytorch loss module, e.g., nn.CrossEntropyLoss().
  validation_metric
      A torchmetrics module used in the validation stage, e.g., torchmetrics.Accuracy().
  validation_metric_name
      Name of validation metric in case that validation_metric is a aggregation metric,
      e.g., torchmetrics.MeanMetric, whose name can't reflect the real metric name.
  custom_metric_func
      A customized metric function in case that torchmetrics doesn't have the metric.
      It is generally used together with torchmetrics' aggregators, e.g., torchmetrics.MeanMetric.
      Refer to https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/aggregation.py
  test_metric
      A torchmetrics module used in the test stage, e.g., torchmetrics.Accuracy().
  track_grad_norm
      Track the p-norm of gradients during training. May be set to ‘inf’ infinity-norm.
      If using Automatic Mixed Precision (AMP), the gradients will be unscaled before logging them.
docstr_tok: ["Parameters","--","--","--","--","--","query_model","The","query","model",".","response_model","The","response","model",".","signature","query","or","response",".","match_label","The","label","of","match","class",".","matches","A","list","of","DictConfigs",",","each","of","which","defines","one","pair","of","feature","column","match","and","the","configs","to","compute","the","matching","loss",".","optim_type","Optimizer","type",".","We","now","support",":","-","adamw","-","adam","-","sgd","lr_choice","How","to","set","each","layer","'s","learning","rate",".","If","not","specified",",","the","default","is","a","single","learnng","rate","for","all","layers",".","Otherwise",",","we","now","support","two","choices",":","-","two_stages","The","layers","in","the","pretrained","models","have","a","small","learning","rate","(","lr","*","lr_mult",")",",","while","the","newly","added","head","layers","use","the","provided","learning","rate",".","-","layerwise_decay","The","layers","have","decreasing","learning","rate","from","the","output","end","to","the","input","end",".","The","intuition","is","that","later","layers","are","more","task-related",",","hence","larger","learning","rates",".","lr_schedule","Learning","rate","schedule",".","We","now","support",":","-","cosine_decay","Linear","warmup","followed","by","cosine","decay","-","polynomial_decay","Linear","warmup","followed","by","polynomial","decay","lr","Learning","rate",".","lr_decay","The","learning","rate","decay","factor","(","0",",","1",")",".","It","is","used","only","when","lr_choice","is","``","layerwise_decay","''",".","end_lr","The","final","learning","rate","after","decay",".","lr_mult","The","learning","rate","multiplier","(","0",",","1",")",".","It","is","used","only","when","lr_choice","is","``","two_stages","''",".","weight_decay","The","weight","decay","to","regularize","layer","weights","'","l2","norm",".","warmup_steps","How","many","steps","to","warmup","learning","rate",".","If","a","float","(","0",",","1",")",",","it","would","represent","the","percentage","of","steps","over","all","the","training","steps",".","The","actual","number","is","calculated","as","``","int","(","warmup_steps","*","max_steps",")","''",".","If","an","integer",",","it","would","be","the","exact","step","number",".","loss_func","A","Pytorch","loss","module",",","e.g.",",","nn.CrossEntropyLoss","(",")",".","validation_metric","A","torchmetrics","module","used","in","the","validation","stage",",","e.g.",",","torchmetrics.Accuracy","(",")",".","validation_metric_name","Name","of","validation","metric","in","case","that","validation_metric","is","a","aggregation","metric",",","e.g.",",","torchmetrics.MeanMetric",",","whose","name","ca","n't","reflect","the","real","metric","name",".","custom_metric_func","A","customized","metric","function","in","case","that","torchmetrics","does","n't","have","the","metric",".","It","is","generally","used","together","with","torchmetrics","'","aggregators",",","e.g.",",","torchmetrics.MeanMetric",".","Refer","to","https",":","\/\/github.com\/PyTorchLightning\/metrics\/blob\/master\/torchmetrics\/aggregation.py","test_metric","A","torchmetrics","module","used","in","the","test","stage",",","e.g.",",","torchmetrics.Accuracy","(",")",".","track_grad_norm","Track","the","p-norm","of","gradients","during","training",".","May","be","set","to","\u2018","inf","\u2019","infinity-norm",".","If","using","Automatic","Mixed","Precision","(","AMP",")",",","the","gradients","will","be","unscaled","before","logging","them","."]
code: def __init__(
      self,
      query_model: nn.Module,
      response_model: nn.Module,
      signature: Optional[str] = None,
      match_label: Optional[int] = None,
      matches: Optional[List[DictConfig]] = None,
      optim_type: Optional[str] = None,
      lr_choice: Optional[str] = None,
      lr_schedule: Optional[str] = None,
      lr: Optional[float] = None,
      lr_decay: Optional[float] = None,
      end_lr: Optional[Union[float, int]] = None,
      lr_mult: Optional[Union[float, int]] = None,
      weight_decay: Optional[float] = None,
      warmup_steps: Optional[int] = None,
      loss_func: Optional[_Loss] = None,
      miner_func: Optional[_Loss] = None,
      validation_metric: Optional[torchmetrics.Metric] = None,
      validation_metric_name: Optional[str] = None,
      custom_metric_func: Callable = None,
      test_metric: Optional[torchmetrics.Metric] = None,
      track_grad_norm: Optional[Union[int, str]] = -1,
  ):
      """
      Parameters
      ----------
      query_model
          The query model.
      response_model
          The response model.
      signature
          query or response.
      match_label
          The label of match class.
      matches
          A list of DictConfigs, each of which defines one pair of feature column match and the configs
          to compute the matching loss.
      optim_type
          Optimizer type. We now support:
          - adamw
          - adam
          - sgd
      lr_choice
          How to set each layer's learning rate. If not specified, the default is a single
          learnng rate for all layers. Otherwise, we now support two choices:
          - two_stages
              The layers in the pretrained models have a small learning rate (lr * lr_mult),
              while the newly added head layers use the provided learning rate.
          - layerwise_decay
              The layers have decreasing learning rate from the output end to the input end.
              The intuition is that later layers are more task-related, hence larger learning rates.
      lr_schedule
          Learning rate schedule. We now support:
          - cosine_decay
              Linear warmup followed by cosine decay
          - polynomial_decay
              Linear warmup followed by polynomial decay
      lr
          Learning rate.
      lr_decay
          The learning rate decay factor (0, 1). It is used only when lr_choice is "layerwise_decay".
      end_lr
          The final learning rate after decay.
      lr_mult
          The learning rate multiplier (0, 1). It is used only when lr_choice is "two_stages".
      weight_decay
          The weight decay to regularize layer weights' l2 norm.
      warmup_steps
          How many steps to warmup learning rate. If a float (0, 1), it would represent the
          percentage of steps over all the training steps. The actual number is calculated as
          "int(warmup_steps * max_steps)". If an integer, it would be the exact step number.
      loss_func
          A Pytorch loss module, e.g., nn.CrossEntropyLoss().
      validation_metric
          A torchmetrics module used in the validation stage, e.g., torchmetrics.Accuracy().
      validation_metric_name
          Name of validation metric in case that validation_metric is a aggregation metric,
          e.g., torchmetrics.MeanMetric, whose name can't reflect the real metric name.
      custom_metric_func
          A customized metric function in case that torchmetrics doesn't have the metric.
          It is generally used together with torchmetrics' aggregators, e.g., torchmetrics.MeanMetric.
          Refer to https://github.com/PyTorchLightning/metrics/blob/master/torchmetrics/aggregation.py
      test_metric
          A torchmetrics module used in the test stage, e.g., torchmetrics.Accuracy().
      track_grad_norm
          Track the p-norm of gradients during training. May be set to ‘inf’ infinity-norm.
          If using Automatic Mixed Precision (AMP), the gradients will be unscaled before logging them.
      """
      super().__init__()
      self.save_hyperparameters(
          ignore=[
              "query_model",
              "response_model",
              "validation_metric",
              "test_metric",
              "loss_func",
              "miner_func",
              "matches",
          ]
      )
      self.query_model = query_model
      self.response_model = response_model
      if signature:
          assert signature in [QUERY, RESPONSE]
      self.signature = signature
      self.validation_metric = validation_metric
      self.validation_metric_name = f"val_{validation_metric_name}"
  
      if (
          isinstance(validation_metric, BaseAggregator)
          and custom_metric_func is None
      ):
          raise ValueError(
              f"validation_metric {validation_metric} is an aggregation metric,"
              f"which must be used with a customized metric function."
          )
      self.custom_metric_func = custom_metric_func
  
      self.matches = matches
      self.match_label = match_label
      self.reverse_prob = match_label == 0
  
      logger.debug(f"match label: {match_label}")
      logger.debug(f"reverse probability: {self.reverse_prob}")
  
      self.loss_func = loss_func
      self.miner_func = miner_func
      self.track_grad_norm = track_grad_norm
  
imports: import logging
  from typing import Callable, Dict, List, Optional, Union
  import lightning.pytorch
  import torch
  import torchmetrics
  from lightning.pytorch.utilities import grad_norm
  from omegaconf import DictConfig
  from torch import nn
  from torch.nn.modules.loss import _Loss
  from torchmetrics.aggregation import BaseAggregator
  from ..constants import FEATURES, LOGIT_SCALE, PROBABILITY, QUERY, RESPONSE
  from ..models.utils import run_model
  from ..utils.matcher import compute_matching_probability
  from .losses import MultiNegativesSoftmaxLoss
  from .utils import CustomHitRate, apply_layerwise_lr_decay, apply_single_lr, apply_two_stages_lr, generate_metric_learning_labels, get_lr_scheduler, get_optimizer
label_desc: Use image node_id 1 to create a new MatcherLitModule object from inherited base classes: pl with example: obj = MatcherLitModule(query_model, response_model, signature, match_label, matches, optim_type, lr_choice, lr_schedule, lr, lr_decay, end_lr, lr_mult, weight_decay, warmup_steps, loss_func, miner_func, validation_metric, validation_metric_name, custom_metric_func, test_metric, track_grad_norm)