class_docstr: use AutoModelForMaskedLM, and select the output by another layer in the lit model
code: def _init_label_word(self):
      args = self.args
      # ./dataset/dataset_name
      dataset_name = args.data_dir.split("/")[1]
      model_name_or_path = args.model_name_or_path.split("/")[-1]
      label_path = f"data/{model_name_or_path}.pt"
      # [num_labels, num_tokens], ignore the unanswerable
      if "dialogue" in args.data_dir:
          label_word_idx = torch.load(label_path)[:-1]
      else:
          label_word_idx = torch.load(label_path)
  
      num_labels = len(label_word_idx)
  
      self.cur_model.resize_token_embeddings(len(self.tokenizer))
      with torch.no_grad():
          word_embeddings = self.cur_model.get_input_embeddings()
          continous_label_word = [
              a[0]
              for a in self.tokenizer(
                  [f"[class{i}]" for i in range(1, num_labels + 1)],
                  add_special_tokens=False,
              )["input_ids"]
          ]
          for i, idx in enumerate(label_word_idx):
              word_embeddings.weight[
                  continous_label_word[i]
              ] = torch.mean(word_embeddings.weight[idx], dim=0)
              # word_embeddings.weight[continous_label_word[i]] = self.relation_embedding[i]
          so_word = [
              a[0]
              for a in self.tokenizer(
                  ["[obj]", "[sub]"], add_special_tokens=False
              )["input_ids"]
          ]
          meaning_word = [
              a[0]
              for a in self.tokenizer(
                  [
                      "person",
                      "organization",
                      "location",
                      "date",
                      "country",
                  ],
                  add_special_tokens=False,
              )["input_ids"]
          ]
  
          for i, idx in enumerate(so_word):
              word_embeddings.weight[so_word[i]] = torch.mean(
                  word_embeddings.weight[meaning_word], dim=0
              )
          assert torch.equal(
              self.cur_model.get_input_embeddings().weight,
              word_embeddings.weight,
          )
          assert torch.equal(
              self.cur_model.get_input_embeddings().weight,
              self.cur_model.get_output_embeddings().weight,
          )
  
      self.word2label = continous_label_word
  
imports: import json
  import torch
  import torch.nn
  import torch.nn.functional
  import numpy
  from .base import BaseLitModel
  from .util import dialog_f1_eval, f1_eval, compute_f1, acc, f1_score
  from transformers.optimization import get_linear_schedule_with_warmup
  from functools import partial
label_desc: Use image node_id 2 for calling the BertLitModel obj's underlying member method code with example usage: obj._init_label_word() without return types