code: def generate_bank_model(args, train_df, predictor):
      bank_keys = []
      bank_labels = []
      with torch.no_grad():
          for augment_idx in range(args.aug_epochs):
              print(
                  "Augment Epoch: {:} / {:}".format(
                      augment_idx, args.aug_epochs
                  )
              )
              train_features = extract_embedding(
                  args, train_df, predictor, args.column_names
              )
              bank_keys.append(train_features.unsqueeze(0))
  
          for index, per_data in train_df.iterrows():
              bank_labels.append(per_data[args.label_column])
  
          bank_keys = torch.cat(bank_keys, dim=0).mean(dim=0)
          bank_keys /= bank_keys.norm(dim=-1, keepdim=True)
          bank_keys = bank_keys.permute(1, 0)
          bank_labels = torch.tensor(bank_labels).cuda()
  
      return bank_keys, bank_labels
  
imports: import torch
  import pandas
  from tqdm import tqdm
  import torch.nn.functional
  import torch.nn
  from autogluon.multimodal import MultiModalPredictor
  from datasets import list_datasets, load_dataset
  from setfit import sample_dataset
  from imagedatasets import build_dataset
  from torch.utils.data import Dataset
label_desc: Use image node_id 6 for calling a global function with example usage: generate_bank_model(args, train_df, predictor) and returns: bank_keys,  bank_labels