code: def create_weighted_mask(labels: torch.Tensor):
      # Check if the tensor is 2D. If not, unsqueeze it to make it 2D
      if len(labels.shape) == 1:
          labels = labels.unsqueeze(0)
  
      weights = torch.zeros_like(labels).float()
      for i in range(labels.shape[0]):
          mask = labels[i] != -100
  
          # Create a tensor to track group ids
          group_ids = torch.zeros_like(labels[i]).int()
          curr_group_id = 0
  
          for j in range(1, len(labels[i])):
              if (
                  mask[j] and not mask[j - 1]
              ):  # switch from masked to unmasked label
                  curr_group_id += 1  # start new group
              group_ids[j] = (
                  curr_group_id if mask[j] else 0
              )  # assign group id if unmasked label
  
          # Count only unmasked labels in each group
          group_counts = torch.bincount(group_ids[mask])
  
          mask_weights = torch.zeros_like(labels[i]).float()
          mask_weights[mask] = 1.0 / group_counts[group_ids[mask]]
  
          weights[i] = mask_weights
  
      return weights.squeeze()
  
imports: import math
  import os
  from contextlib import contextmanager
  from functools import partial
  from typing import List
  import numpy
  import torch
  import torch.cuda
  from accelerate.logging import get_logger
  from datasets import set_caching_enabled
  from torch.utils.data import DataLoader, RandomSampler
  from axolotl.core.trainer_builder import HFCausalTrainerBuilder, HFDPOTrainerBuilder
  from axolotl.utils.distributed import is_main_process, reduce_and_broadcast, zero_first
  from axolotl.utils.samplers import MultipackBatchSampler, get_dataset_lengths
label_desc: Use image node_id 2 for calling a global function with example usage: create_weighted_mask(labels) and returns: weights