Дистилляция BERT для задачи классификации

import re import pandas as pd from pymystem3 import Mystem  # get data data = pd.read_csv('data.csv')  texts = list(data['comment']) labels = list(map(int, data['toxic'].values))  # clean texts texts = [re.sub('[^а-яё ]', ' ', str(t).lower()) for t in texts] texts = [re.sub(r" +", " ", t).strip() for t in texts]  # lemmatize mstm = Mystem()  normalized = [''.join(mstm.lemmatize(t)[:-1]) for t in texts]  # remove stopwords with open('./stopwords.txt') as f:     stopwords = [line.rstrip('\n') for line in f]  def drop_stop(text):     tokens = text.split(' ')     tokens = [t for t in tokens if t not in stopwords]     return ' '.join(tokens)  normalized = [drop_stop(text) for text in normalized]  # new dataset df = pd.DataFrame() df['text'] = texts df['norm'] = normalized df['label'] = labels  # train-valid-test-split from sklearn.model_selection import train_test_split  train, test = train_test_split(df, test_size=0.3, random_state=42) valid, test = train_test_split(test, test_size=0.5, random_state=42)  # tf-idf from sklearn.feature_extraction.text import TfidfVectorizer  model_tfidf = TfidfVectorizer(max_features=5000)  train_tfidf = model_tfidf.fit_transform(train['norm'].values) valid_tfidf = model_tfidf.transform(valid['norm'].values) test_tfidf = model_tfidf.transform(test['norm'].values)  # RF from sklearn.ensemble import RandomForestClassifier  cls = RandomForestClassifier(random_state=42) cls.fit(train_tfidf, train['label'].values)  # prediction predictions = cls.predict(test_tfidf)  # score from sklearn.metrics import f1_score  f1_score(predictions, test['label'].values)

# get data import pandas as pd  train = pd.read_csv('train.csv') valid = pd.read_csv('valid.csv') test = pd.read_csv('test.csv')  # create tokenizer from tokenizers import Tokenizer from tokenizers import ByteLevelBPETokenizer from tokenizers.pre_tokenizers import Whitespace  tokenizer = ByteLevelBPETokenizer() tokenizer.pre_tokenizer = Whitespace() tokenizer.enable_padding(pad_id=0, pad_token='<pad>')  texts_path = 'texts.txt'  with open(texts_path, 'w') as f:     for text in list(train['text'].values):         f.write("%s\n" % text)  tokenizer.train(     files=[texts_path],     vocab_size=5_000,     min_frequency=2,     special_tokens=['<pad>', '<unk>']     )  # create dataset import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader  class CustomDataset(Dataset):      def __init__(self, tokens, labels, max_len):         self.tokens = tokens         self.labels = labels         self.max_len = max_len       def __len__(self):         return len(self.tokens)       def __getitem__(self, idx):         label = self.labels[idx]         label = torch.tensor(label)         tokens = self.tokens[idx]         out = torch.zeros(self.max_len, dtype=torch.long)         out[:len(tokens)] = torch.tensor(tokens, dtype=torch.long)[:self.max_len]         return out, label  max_len = 64 BATCH_SIZE = 16  train_labels = list(train['label']) train_tokens = [tokenizer.encode(text).ids for text in list(train['text'])] train_dataset = CustomDataset(train_tokens, train_labels, max_len) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False)  test_labels = list(test['label']) test_tokens = [tokenizer.encode(text).ids for text in list(test['text'])] test_dataset = CustomDataset(test_tokens, test_labels, max_len) test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)  # create BiLSTM class LSTM_classifier(nn.Module):       def __init__(self, hidden_dim=128, vocab_size=5000, embedding_dim=300, linear_dim=128, dropout=0.3, n_classes=2):         super().__init__()         self.embedding_layer = nn.Embedding(vocab_size, embedding_dim)         self.lstm_layer = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, bidirectional=True)         self.dropout_layer = nn.Dropout(dropout)                 self.fc_layer = nn.Linear(hidden_dim * 2, linear_dim)         self.batchnorm = nn.BatchNorm1d(linear_dim)         self.relu = nn.ReLU()         self.out_layer = nn.Linear(linear_dim, n_classes)       def forward(self, inputs):         batch_size = inputs.size(0)         embeddings = self.embedding_layer(inputs)         lstm_out, (ht, ct) = self.lstm_layer(embeddings)         out = ht.transpose(0, 1)         out = out.reshape(batch_size, -1)         out = self.fc_layer(out)         out = self.batchnorm(out)         out = self.relu(out)         out = self.dropout_layer(out)         out = self.out_layer(out)         out = torch.squeeze(out, 1)         out = torch.sigmoid(out)         return out  def init_weights(m):     if isinstance(m, nn.Linear):         torch.nn.init.xavier_uniform_(m.weight)         m.bias.data.fill_(0.01)  def eval_nn(model, data_loader):     predicted = []     labels = []     model.eval()     with torch.no_grad():         for data in data_loader:             x, y = data             x = x.to(device)             outputs = model(x)             _, predict = torch.max(outputs.data, 1)             predict = predict.cpu().detach().numpy().tolist()             predicted += predict             labels += y         score = f1_score(labels, predicted, average='binary')     return score  def train_nn(model, optimizer, loss_function, train_loader, test_loader, device, epochs=20):     best_score = 0     for epoch in range(epochs):         model.train()         for inputs, labels in tqdm(train_loader):             inputs, labels = inputs.to(device), labels.to(device)             optimizer.zero_grad()             predict = model(inputs)             loss = loss_function(predict, labels)             loss.backward()             optimizer.step()        score = eval_nn(model, test_loader)         print(epoch, 'valid:', score)         if score > best_score:             torch.save(model.state_dict(),'lstm.pt')             best_score = score     return best_score  # fit NN device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  model = LSTM_classifier(hidden_dim=256, vocab_size=5000, embedding_dim=300, linear_dim=128, dropout=0.1)     model.apply(init_weights)  model.to(device)  optimizer = optim.AdamW(model.parameters())  loss_function = nn.CrossEntropyLoss().to(device)  train_nn(model, optimizer, loss_function, train_loader, valid_loader, device, epochs=20)  eval_nn(model, test_loader)

import torch from torch.utils.data import Dataset  class BertDataset(Dataset):    def __init__(self, texts, targets, tokenizer, max_len=512):     self.texts = texts     self.targets = targets     self.tokenizer = tokenizer     self.max_len = max_len    def __len__(self):     return len(self.texts)    def __getitem__(self, idx):     text = str(self.texts[idx])     target = self.targets[idx]      encoding = self.tokenizer.encode_plus(         text,         add_special_tokens=True,         max_length=self.max_len,         return_token_type_ids=False,         padding='max_length',         return_attention_mask=True,         return_tensors='pt',         truncation=True     )      return {       'text': text,       'input_ids': encoding['input_ids'].flatten(),       'attention_mask': encoding['attention_mask'].flatten(),       'targets': torch.tensor(target, dtype=torch.long)     }  from tqdm import tqdm import numpy as np import torch from transformers import BertTokenizer, BertForSequenceClassification from torch.utils.data import Dataset, DataLoader from transformers import AdamW, get_linear_schedule_with_warmup from sklearn.metrics import precision_recall_fscore_support   class BertClassifier:      def __init__(self, path, n_classes=2):         self.path = path         self.model = BertForSequenceClassification.from_pretrained(path)         self.tokenizer = BertTokenizer.from_pretrained(path)         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")         self.max_len = 512         self.out_features = self.model.bert.encoder.layer[1].output.dense.out_features         self.model.classifier = torch.nn.Linear(self.out_features, n_classes)         self.model.to(self.device)           def preparation(self, X_train, y_train, epochs):         # create datasets         self.train_set = BertDataset(X_train, y_train, self.tokenizer)         # create data loaders         self.train_loader = DataLoader(self.train_set, batch_size=2, shuffle=True)         # helpers initialization         self.optimizer = AdamW(             self.model.parameters(),             lr=2e-5,             weight_decay=0.005,             correct_bias=True             )         self.scheduler = get_linear_schedule_with_warmup(                 self.optimizer,                 num_warmup_steps=500,                 num_training_steps=len(self.train_loader) * epochs             )         self.loss_fn = torch.nn.CrossEntropyLoss().to(self.device)       def fit(self):         self.model = self.model.train()         losses = []         correct_predictions = 0          for data in tqdm(self.train_loader):             input_ids = data["input_ids"].to(self.device)             attention_mask = data["attention_mask"].to(self.device)             targets = data["targets"].to(self.device)              outputs = self.model(                 input_ids=input_ids,                 attention_mask=attention_mask                 )              preds = torch.argmax(outputs.logits, dim=1)             loss = self.loss_fn(outputs.logits, targets)              correct_predictions += torch.sum(preds == targets)              losses.append(loss.item())                          loss.backward()             torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)             self.optimizer.step()             self.scheduler.step()             self.optimizer.zero_grad()          train_acc = correct_predictions.double() / len(self.train_set)         train_loss = np.mean(losses)         return train_acc, train_loss           def train(self, X_train, y_train, X_valid, y_valid, X_test, y_test, epochs=1):         print('*' * 10)         print(f'Model: {self.path}')         self.preparation(X_train, y_train, epochs)         for epoch in range(epochs):             print(f'Epoch {epoch + 1}/{epochs}')             train_acc, train_loss = self.fit()             print(f'Train loss {train_loss} accuracy {train_acc}')             predictions_valid = [self.predict(x) for x in X_valid]             precision, recall, f1score = precision_recall_fscore_support(y_valid, predictions_valid, average='macro')[:3]             print('Valid:')             print(f'precision: {precision}, recall: {recall}, f1score: {f1score}')             predictions_test = [self.predict(x) for x in X_test]             precision, recall, f1score = precision_recall_fscore_support(y_test, predictions_test, average='macro')[:3]             print('Test:')             print(f'precision: {precision}, recall: {recall}, f1score: {f1score}')         print('*' * 10)          def predict(self, text):         self.model = self.model.eval()         encoding = self.tokenizer.encode_plus(             text,             add_special_tokens=True,             max_length=self.max_len,             return_token_type_ids=False,             truncation=True,             padding='max_length',             return_attention_mask=True,             return_tensors='pt',         )                  out = {               'text': text,               'input_ids': encoding['input_ids'].flatten(),               'attention_mask': encoding['attention_mask'].flatten()           }                  input_ids = out["input_ids"].to(self.device)         attention_mask = out["attention_mask"].to(self.device)                  outputs = self.model(             input_ids=input_ids.unsqueeze(0),             attention_mask=attention_mask.unsqueeze(0)         )                  prediction = torch.argmax(outputs.logits, dim=1).cpu().numpy()[0]          return prediction  import pandas as pd  train = pd.read_csv('train.csv') valid = pd.read_csv('valid.csv') test = pd.read_csv('test.csv')  classifier = BertClassifier(     path='cointegrated/rubert-tiny',     n_classes=2 )  classifier.train(         X_train=list(train['text']),         y_train=list(train['label']),         X_valid=list(valid['text']),         y_valid=list(valid['label']),         X_test=list(test['text']),         y_test=list(test['label']),         epochs=1 )  path = './trainer' classifier.model.save_pretrained(path) classifier.tokenizer.save_pretrained(path)

import torch import torch.nn as nn import torch.optim as optim  from torch.utils.data import Dataset, DataLoader  from tokenizers import Tokenizer from tokenizers import ByteLevelBPETokenizer from tokenizers.pre_tokenizers import Whitespace  from transformers import BertTokenizer, BertForSequenceClassification from transformers import AdamW, get_linear_schedule_with_warmup  from sklearn.metrics import precision_recall_fscore_support from sklearn.metrics import f1_score  import numpy as np import pandas as pd  ### data  train = pd.read_csv('train.csv') test = pd.read_csv('test.csv')  ### tokenizer: train  tokenizer = ByteLevelBPETokenizer() tokenizer.pre_tokenizer = Whitespace() tokenizer.enable_padding(pad_id=0, pad_token='<pad>')  texts_path = 'texts.txt'  with open(texts_path, 'w') as f:     for text in list(train['text'].values):         f.write("%s\n" % text)  tokenizer.train(     files=[texts_path],     vocab_size=5_000,     min_frequency=2,     special_tokens=['<pad>', '<unk>']     )  ### load BERT tokenizer  tokenizer_bert = BertTokenizer.from_pretrained('./rubert-tiny')  ### dataset  class CustomDataset(Dataset):      def __init__(self, tokens, labels, max_len):         self.tokens = tokens         self.labels = labels         self.max_len = max_len       def __len__(self):         return len(self.tokens)       def __getitem__(self, idx):         label = self.labels[idx]         label = torch.tensor(label)         tokens = self.tokens[idx]         out = torch.zeros(self.max_len, dtype=torch.long)         out[:len(tokens)] = torch.tensor(tokens, dtype=torch.long)[:self.max_len]         return out, label  max_len = 64 BATCH_SIZE = 16  train_labels = list(train['label']) train_tokens = [tokenizer.encode(str(text)).ids for text in list(train['text'])] train_dataset = CustomDataset(train_tokens, train_labels, max_len) train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=False)  test_labels = list(test['label']) test_tokens = [tokenizer.encode(str(text)).ids for text in list(test['text'])] test_dataset = CustomDataset(test_tokens, test_labels, max_len) test_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)  class LSTM_classifier(nn.Module):       def __init__(self, hidden_dim=128, vocab_size=5000, embedding_dim=300, linear_dim=128, dropout=0.3, n_classes=2):         super().__init__()         self.embedding_layer = nn.Embedding(vocab_size, embedding_dim)         self.lstm_layer = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, bidirectional=True)         self.dropout_layer = nn.Dropout(dropout)                 self.fc_layer = nn.Linear(hidden_dim * 2, linear_dim)         self.batchnorm = nn.BatchNorm1d(linear_dim)         self.relu = nn.ReLU()         self.out_layer = nn.Linear(linear_dim, n_classes)       def forward(self, inputs):         batch_size = inputs.size(0)         embeddings = self.embedding_layer(inputs)         lstm_out, (ht, ct) = self.lstm_layer(embeddings)         out = ht.transpose(0, 1)         out = out.reshape(batch_size, -1)         out = self.fc_layer(out)         out = self.batchnorm(out)         out = self.relu(out)         out = self.dropout_layer(out)         out = self.out_layer(out)         out = torch.squeeze(out, 1)         out = torch.sigmoid(out)         return out  ########  def init_weights(m):     if isinstance(m, nn.Linear):         torch.nn.init.xavier_uniform_(m.weight)         m.bias.data.fill_(0.01)  def eval_nn(model, data_loader):     predicted = []     labels = []     model.eval()     with torch.no_grad():         for data in data_loader:             x, y = data             x = x.to(device)             outputs = model(x)             _, predict = torch.max(outputs.data, 1)             predict = predict.cpu().detach().numpy().tolist()             predicted += predict             labels += y         score = f1_score(labels, predicted, average='binary')     return score  def train_nn(model, optimizer, loss_function, train_loader, test_loader, device, epochs=20):     best_score = 0     for epoch in range(epochs):         model.train()         for inputs, labels in train_loader:             inputs, labels = inputs.to(device), labels.to(device)             optimizer.zero_grad()             predict = model(inputs)             loss = loss_function(predict, labels)             loss.backward()             optimizer.step()         score = eval_nn(model, test_loader)         print(epoch, 'valid:', score)         if score > best_score:             torch.save(model.state_dict(), 'lstm.pt')             best_score = score     return best_score  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  model = LSTM_classifier(hidden_dim=256, vocab_size=5000, embedding_dim=300, linear_dim=128, dropout=0.1)     model.apply(init_weights);  model.to(device);  optimizer = optim.AdamW(model.parameters())  loss_function = nn.CrossEntropyLoss().to(device)  train_nn(model, optimizer, loss_function, train_loader, test_loader, device, epochs=3)  #####  class DistillDataset(Dataset):      def __init__(self, texts, labels, tokenizer_bert, tokenizer_lstm, max_len):         self.texts = texts         self.labels = labels         self.tokenizer_bert = tokenizer_bert         self.tokenizer_lstm = tokenizer_lstm         self.max_len = max_len       def __len__(self):         return len(self.texts)       def __getitem__(self, idx):         text = self.texts[idx]         label = self.labels[idx]         label = torch.tensor(label)         # lstm         tokens_lstm = self.tokenizer_lstm.encode(str(text)).ids         out_lstm = torch.zeros(self.max_len, dtype=torch.long)         out_lstm[:len(tokens_lstm)] = torch.tensor(tokens_lstm, dtype=torch.long)[:self.max_len]         # bert         encoding = self.tokenizer_bert.encode_plus(             str(text),             add_special_tokens=True,             max_length=self.max_len,             return_token_type_ids=False,             truncation=True,             padding='max_length',             return_attention_mask=True,             return_tensors='pt',         )                  out_bert = {               'input_ids': encoding['input_ids'].flatten(),               'attention_mask': encoding['attention_mask'].flatten()         }         return out_lstm, out_bert, label  train_dataset_distill = DistillDataset(     list(train['text']),     list(train['label']),     tokenizer_bert,     tokenizer,     max_len )  train_loader_distill = DataLoader(train_dataset_distill, batch_size=BATCH_SIZE, shuffle=True)  ### BERT-teacher model  class BertTrainer:      def __init__(self, path_model, n_classes=2):         self.model = BertForSequenceClassification.from_pretrained(path_model, num_labels=n_classes)         self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")         self.max_len = 512         self.model.to(self.device)         self.model = self.model.eval()          def predict(self, inputs):              input_ids = inputs["input_ids"].to(self.device)         attention_mask = inputs["attention_mask"].to(self.device)         with torch.no_grad():             outputs = self.model(                 input_ids=input_ids,                 attention_mask=attention_mask             )         return outputs.logits  teacher = BertTrainer('./rubert-tiny')  ### BiLSTM-student model  class CustomLSTM(nn.Module):       def __init__(self, hidden_dim=128, vocab_size=5000, embedding_dim=300, linear_dim=128, dropout=0.3, n_classes=2):         super().__init__()         self.embedding_layer = nn.Embedding(vocab_size, embedding_dim)         self.lstm_layer = nn.LSTM(embedding_dim, hidden_dim, batch_first=True, bidirectional=True)         self.dropout_layer = nn.Dropout(dropout)                 self.fc_layer = nn.Linear(hidden_dim * 2, linear_dim)         self.batchnorm = nn.BatchNorm1d(linear_dim)         self.relu = nn.ReLU()         self.out_layer = nn.Linear(linear_dim, n_classes)       def forward(self, inputs):         batch_size = inputs.size(0)         embeddings = self.embedding_layer(inputs)         lstm_out, (ht, ct) = self.lstm_layer(embeddings)         out = ht.transpose(0, 1)         out = out.reshape(batch_size, -1)         out = self.fc_layer(out)         out = self.batchnorm(out)         out = self.relu(out)         out = self.dropout_layer(out)         out = self.out_layer(out) #         out = torch.squeeze(out, 1) #         out = torch.sigmoid(out)         return out  def loss_function(output, teacher_prob, real_label, a=0.5):     criterion_mse = torch.nn.MSELoss()     criterion_ce = torch.nn.CrossEntropyLoss()     return a * criterion_ce(output, real_label) + (1 - a) * criterion_mse(output, teacher_prob)  def init_weights(m):     if isinstance(m, nn.Linear):         torch.nn.init.xavier_uniform_(m.weight)         m.bias.data.fill_(0.01)  def eval_nn(model, data_loader):     predicted = []     labels = []     model.eval()     with torch.no_grad():         for data in data_loader:             x, y = data             x = x.to(device)             outputs = model(x)             _, predict = torch.max(outputs.data, 1)             predict = predict.cpu().detach().numpy().tolist()             predicted += predict             labels += y         score = f1_score(labels, predicted, average='binary')     return labels, predicted, score  def train_distill(model, teacher, optimizer, loss_function, distill_loader, train_loader, test_loader, device, epochs=30, alpha=0.5):     best_score = 0     score_list = []     for epoch in range(epochs):         model.train()         for inputs, inputs_teacher, labels in distill_loader:             inputs, labels = inputs.to(device), labels.to(device)             optimizer.zero_grad()             predict = model(inputs)             teacher_predict = teacher.predict(inputs_teacher)             loss = loss_function(predict, teacher_predict, labels, alpha)             loss.backward()             optimizer.step()         score_train = round(eval_nn(model, train_loader)[2], 3)         score_test = round(eval_nn(model, test_loader)[2], 3)         score_list.append((score_train, score_test))         print(epoch, score_train, score_test)         if score_test > best_score:                  best_score = score_test             best_model = model     torch.save(best_model.state_dict(), f'./results/lstm_{best_score}.pt')     return best_model, best_score, score_list  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")  vocab_size = tokenizer.get_vocab_size()  vocab_size  score_alpha = [] for alpha in [0, 0.25, 0.5, 0.75, 1]:     model = LSTM_classifier(hidden_dim=256, vocab_size=5000, embedding_dim=300, linear_dim=128, dropout=0.1)         model.apply(init_weights)     model.to(device)     optimizer = optim.AdamW(model.parameters())     _, _, score_list = train_distill(model, teacher, optimizer, loss_function, train_loader_distill, train_loader, test_loader, device, 30, alpha)     score_alpha.append(score_list)  import matplotlib.pyplot as plt import numpy as np  a_list = [1, 0.75, 0.5, 0.25, 0]  for i, score in enumerate(score_alpha):     _, score_test = list(zip(*score))     plt.plot(score_test, label=f'{a_list[i]}') plt.grid(True) plt.legend() plt.show()

from transformers import BertTokenizerFast, BertForPreTraining, BertModel, BertConfig from collections import Counter from tqdm.auto import tqdm, trange import pandas as pd  train = pd.read_csv('train.csv') X_train=list(train['text'])  tokenizer = BertTokenizerFast.from_pretrained('./rubert-tiny')  cnt = Counter() for text in tqdm(X_train):     cnt.update(tokenizer(str(text))['input_ids'])  resulting_vocab = {     tokenizer.vocab[k] for k in tokenizer.special_tokens_map.values() }  for k, v in cnt.items():     if v > 5:         resulting_vocab.add(k)  resulting_vocab = sorted(resulting_vocab)  tokenizer.save_pretrained('./bert_distill');  inv_voc = {idx: word for word, idx in tokenizer.vocab.items()}  with open('./bert_distill/vocab.txt', 'w', encoding='utf-8') as f:     for idx in resulting_vocab:         f.write(inv_voc[idx] + '\n')

config = BertConfig(     emb_size=256,     hidden_size=256,     intermediate_size=256,     max_position_embeddings=512,     num_attention_heads=8,     num_hidden_layers=3,     vocab_size=tokenizer_distill.vocab_size )  model = BertForPreTraining(config)  model.save_pretrained('./bert_distill')  from transformers import BertModel # load model without CLS-head teacher = BertForPreTraining.from_pretrained('./rubert-tiny')  tokenizer_teacher = BertTokenizerFast.from_pretrained('./rubert-tiny')  # copy input embeddings accordingly with resulting_vocab model.bert.embeddings.word_embeddings.weight.data = teacher.bert.embeddings.word_embeddings.weight.data[resulting_vocab, :256].clone() model.bert.embeddings.position_embeddings.weight.data = teacher.bert.embeddings.position_embeddings.weight.data[:, :256].clone()  # copy output embeddings model.cls.predictions.decoder.weight.data = teacher.cls.predictions.decoder.weight.data[resulting_vocab, :256].clone()

inputs = tokenizer_distill(texts, return_tensors='pt', padding=True, truncation=True, max_length=16) inputs = preprocess_inputs(inputs, tokenizer_distill, data_collator) outputs = model(**inputs, output_hidden_states=True) loss += nn.CrossEntropyLoss(         outputs.prediction_logits.view(-1, model.config.vocab_size),         inputs['labels'].view(-1)     )

def loss_kl(inputs, outputs, model, teacher, vocab_mapping, temperature=1.0):     new_inputs = torch.tensor(         [[vocab_mapping[i] for i in row] for row in inputs['input_ids']]     ).to(inputs['input_ids'].device)     with torch.no_grad():         teacher_out = teacher(             input_ids=new_inputs,              token_type_ids=inputs['token_type_ids'],             attention_mask=inputs['attention_mask']         )     # the whole batch, all tokens after the [cls], the whole dimension     kd_loss = torch.nn.KLDivLoss(reduction='batchmean')(         F.log_softmax(outputs.prediction_logits[:, 1:, :] / temperature, dim=1),          F.softmax(teacher_out.prediction_logits[:, 1:, vocab_mapping] / temperature, dim=1)     ) / outputs.prediction_logits.shape[-1]     return kd_loss

input_teacher = {k: v for k, v in tokenizer_teacher(         texts,         return_tensors='pt',         padding=True,         max_length=16,         truncation=True     ).items()}  with torch.no_grad():     out_teacher = teacher_mse(**input_teacher)  embeddings_teacher_norm = torch.nn.functional.normalize(out_teacher.pooler_output)  input_distill = {k: v for k, v in tokenizer_distill(         texts,         return_tensors='pt',         padding=True,         max_length=16,         truncation=True     ).items()}  out = model(**input_distill, output_hidden_states=True) embeddings = model.bert.pooler(out.hidden_states[-1]) embeddings_norm = torch.nn.functional.normalize(adapter_emb(embeddings)) loss = torch.nn.MSELoss(embeddings_norm, embeddings_teacher_norm)