import contextlib import torch from model.blip2qformer import Blip2Qformer import pytorch_lightning as pl from torch import optim from lavis.common.optims import LinearWarmupCosineLRScheduler, LinearWarmupStepLRScheduler from tqdm import tqdm from model.help_funcs import AttrDict, pad_and_concat from typing import Any, Dict class Blip2Stage1(pl.LightningModule): def __init__(self, args): super().__init__() if isinstance(args, dict): args = AttrDict(**args) self.args = args self.rerank_cand_num = args.rerank_cand_num self.blip2qformer = Blip2Qformer(args.ptm, args.lm, args.bert_name, args.plm_name, args.temperature, args.plm_tune, args.num_query_token, args.cross_attention_freq, args.projection_dim, args.pool_size, args.load_4bit) self.save_hyperparameters(args) def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None: # checkpoint.pop('optimizer_states') to_be_removed = [] for key, value in checkpoint['state_dict'].items(): try: if not self.get_parameter(key).requires_grad: to_be_removed.append(key) except AttributeError: to_be_removed.append(key) for key in to_be_removed: checkpoint['state_dict'].pop(key) def maybe_autocast(self, dtype=torch.bfloat16): # if on cpu, don't use autocast # if on gpu, use autocast with dtype if provided, otherwise use torch.float16 enable_autocast = self.device != torch.device("cpu") if enable_autocast: return torch.cuda.amp.autocast(dtype=dtype) else: return contextlib.nullcontext() def configure_optimizers(self): self.trainer.fit_loop.setup_data() warmup_steps = min(len(self.trainer.train_dataloader), self.args.warmup_steps) optimizer = optim.AdamW(self.parameters(), lr=self.args.init_lr, weight_decay=self.args.weight_decay) if self.args.scheduler == 'linear_warmup_cosine_lr': self.scheduler = LinearWarmupCosineLRScheduler(optimizer, self.args.max_epochs, self.args.min_lr, self.args.init_lr, warmup_steps, self.args.warmup_lr) elif self.args.scheduler == 'linear_warmup_step_lr': self.scheduler = LinearWarmupStepLRScheduler(optimizer, self.args.max_epochs, self.args.min_lr, self.args.init_lr, self.args.lr_decay_rate, self.args.warmup_lr, warmup_steps) elif self.args.scheduler == 'None': self.scheduler = None else: raise NotImplementedError() return optimizer @torch.no_grad() def validation_step(self, batch, batch_idx, dataloader_idx=0): prot_batch, text_batch = batch batch_size = prot_batch.input_ids.shape[0] blip2_loss = self.blip2qformer(batch) ###============== Overall Loss ===================### self.log(f"loader{dataloader_idx}/val_loss_ptc", float(blip2_loss.loss_itc), batch_size=batch_size, sync_dist=True) self.log(f"loader{dataloader_idx}/val_loss_ptm", float(blip2_loss.loss_itm), batch_size=batch_size, sync_dist=True) self.log(f"loader{dataloader_idx}/val_loss_lm", float(blip2_loss.loss_lm), batch_size=batch_size, sync_dist=True) self.log(f"loader{dataloader_idx}/val_loss", float(blip2_loss.loss), batch_size=batch_size, sync_dist=True) return blip2_loss.loss def get_precision(self, precision): if precision in {'16', '16-mixed'}: return torch.float16 elif precision in {'bf16', 'bf16-mixed'}: return torch.bfloat16 elif precision in {'32',}: return torch.float32 else: raise NotImplementedError def retrieval_evaluation_and_log(self, match_dataloader, log_prefix="") -> None: with self.maybe_autocast(self.get_precision(self.trainer.precision)): ## for onto test set p2t_acc, t2p_acc, p2t_rec20, t2p_rec20, \ p2t_rerank_acc, t2p_rerank_acc, p2t_rerank_rec20, t2p_rerank_rec20, \ prot_feat_total, text_feat_total, prot_embed_total, prot_mask_total, text_total, text_mask_total = \ eval_retrieval_inbatch_with_rerank(self.blip2qformer, match_dataloader, self.device) self.log(f"{log_prefix}inbatch_p2t_acc", p2t_acc, sync_dist=False) self.log(f"{log_prefix}inbatch_t2p_acc", t2p_acc, sync_dist=False) self.log(f"{log_prefix}inbatch_p2t_rec20", p2t_rec20, sync_dist=False) self.log(f"{log_prefix}inbatch_t2p_rec20", t2p_rec20, sync_dist=False) self.log(f"{log_prefix}rerank_inbatch_p2t_acc", p2t_rerank_acc, sync_dist=False) self.log(f"{log_prefix}rerank_inbatch_t2p_acc", t2p_rerank_acc, sync_dist=False) self.log(f"{log_prefix}rerank_inbatch_p2t_rec20", p2t_rerank_rec20, sync_dist=False) self.log(f"{log_prefix}rerank_inbatch_t2p_rec20", t2p_rerank_rec20, sync_dist=False) p2t_acc, p2t_rec20, t2p_acc, t2p_rec20, sim_p2t = \ eval_retrieval_fullset(prot_feat_total, text_feat_total, self.device) self.log(f"{log_prefix}fullset_p2t_acc", p2t_acc, sync_dist=False) self.log(f"{log_prefix}fullset_t2p_acc", t2p_acc, sync_dist=False) self.log(f"{log_prefix}fullset_p2t_rec20", p2t_rec20, sync_dist=False) self.log(f"{log_prefix}fullset_t2p_rec20", t2p_rec20, sync_dist=False) p2t_acc, p2t_rec20, t2p_acc, t2p_rec20 = \ eval_retrieval_fullset_for_rerank(self.blip2qformer, sim_p2t, prot_embed_total, prot_mask_total, text_total, text_mask_total, self.rerank_cand_num, self.device) self.log(f"{log_prefix}rerank_fullset_p2t_acc", p2t_acc, sync_dist=False) self.log(f"{log_prefix}rerank_fullset_t2p_acc", t2p_acc, sync_dist=False) self.log(f"{log_prefix}rerank_fullset_p2t_rec20", p2t_rec20, sync_dist=False) self.log(f"{log_prefix}rerank_fullset_t2p_rec20", t2p_rec20, sync_dist=False) def on_validation_epoch_end(self) -> None: if self.current_epoch == 0 or (self.current_epoch + 1) % self.args.retrieval_eval_epoch != 0: return if self.trainer.global_rank == 0: # 可在此处添加其他非评估逻辑（如自定义日志记录） pass # if self.current_epoch == 0 or (self.current_epoch + 1) % self.args.retrieval_eval_epoch != 0: # return # if self.trainer.global_rank == 0: # ## evaluation for mix dataloaders # if hasattr(self, 'swiss_test_match_loader') and hasattr(self, 'onto_test_match_loader'): # self.retrieval_evaluation_and_log(self.swiss_test_match_loader, log_prefix="swiss_test_") # self.retrieval_evaluation_and_log(self.onto_test_match_loader, log_prefix="onto_test_") # return # with self.maybe_autocast(self.get_precision(self.trainer.precision)): # ## for validation set # p2t_acc, t2p_acc, p2t_rec20, t2p_rec20, \ # p2t_rerank_acc, t2p_rerank_acc, p2t_rerank_rec20, t2p_rerank_rec20,\ # prot_feat_total, text_feat_total, _, _, _, _ = \ # eval_retrieval_inbatch_with_rerank(self.blip2qformer, self.val_match_loader, self.device) # self.log("val_inbatch_p2t_acc", p2t_acc, sync_dist=False) # self.log("val_inbatch_t2p_acc", t2p_acc, sync_dist=False) # self.log("val_inbatch_p2t_rec20", p2t_rec20, sync_dist=False) # self.log("val_inbatch_t2p_rec20", t2p_rec20, sync_dist=False) # self.log("rerank_val_inbatch_p2t_acc", p2t_rerank_acc, sync_dist=False) # self.log("rerank_val_inbatch_t2p_acc", t2p_rerank_acc, sync_dist=False) # self.log("rerank_val_inbatch_p2t_rec20", p2t_rerank_rec20, sync_dist=False) # self.log("rerank_val_inbatch_t2p_rec20", t2p_rerank_rec20, sync_dist=False) # p2t_acc, p2t_rec20, t2p_acc, t2p_rec20, _ = \ # eval_retrieval_fullset(prot_feat_total, text_feat_total, self.device) # self.log("val_fullset_p2t_acc", p2t_acc, sync_dist=False) # self.log("val_fullset_t2p_acc", t2p_acc, sync_dist=False) # self.log("val_fullset_p2t_rec20", p2t_rec20, sync_dist=False) # self.log("val_fullset_t2p_rec20", t2p_rec20, sync_dist=False) # ## for test set # p2t_acc, t2p_acc, p2t_rec20, t2p_rec20, \ # p2t_rerank_acc, t2p_rerank_acc, p2t_rerank_rec20, t2p_rerank_rec20, \ # prot_feat_total, text_feat_total, prot_embed_total, prot_mask_total, text_total, text_mask_total = \ # eval_retrieval_inbatch_with_rerank(self.blip2qformer, self.test_match_loader, self.device) # self.log("rerank_test_inbatch_p2t_acc", p2t_rerank_acc, sync_dist=False) # self.log("rerank_test_inbatch_t2p_acc", t2p_rerank_acc, sync_dist=False) # self.log("rerank_test_inbatch_p2t_rec20", p2t_rerank_rec20, sync_dist=False) # self.log("rerank_test_inbatch_t2p_rec20", t2p_rerank_rec20, sync_dist=False) # self.log("test_inbatch_p2t_acc", p2t_acc, sync_dist=False) # self.log("test_inbatch_t2p_acc", t2p_acc, sync_dist=False) # self.log("test_inbatch_p2t_rec20", p2t_rec20, sync_dist=False) # self.log("test_inbatch_t2p_rec20", t2p_rec20, sync_dist=False) # p2t_acc, p2t_rec20, t2p_acc, t2p_rec20, sim_p2t = \ # eval_retrieval_fullset(prot_feat_total, text_feat_total, self.device) # self.log("test_fullset_p2t_acc", p2t_acc, sync_dist=False) # self.log("test_fullset_t2p_acc", t2p_acc, sync_dist=False) # self.log("test_fullset_p2t_rec20", p2t_rec20, sync_dist=False) # self.log("test_fullset_t2p_rec20", t2p_rec20, sync_dist=False) # p2t_acc, p2t_rec20, t2p_acc, t2p_rec20 = \ # eval_retrieval_fullset_for_rerank(self.blip2qformer, sim_p2t, prot_embed_total, prot_mask_total, text_total, text_mask_total, self.rerank_cand_num, self.device) # self.log("rerank_test_fullset_p2t_acc", p2t_acc, sync_dist=False) # self.log("rerank_test_fullset_t2p_acc", t2p_acc, sync_dist=False) # self.log("rerank_test_fullset_p2t_rec20", p2t_rec20, sync_dist=False) # self.log("rerank_test_fullset_t2p_rec20", t2p_rec20, sync_dist=False) # del prot_feat_total, text_feat_total def training_step(self, batch, batch_idx): self.scheduler.step(self.trainer.current_epoch, self.trainer.global_step) prot_batch, text_batch = batch batch_size = prot_batch.input_ids.shape[0] blip2_loss = self.blip2qformer(batch) ###============== Overall Loss ===================### self.log("train_loss_ptc", float(blip2_loss.loss_itc), batch_size=batch_size, sync_dist=True) self.log("train_loss_ptm", float(blip2_loss.loss_itm), batch_size=batch_size, sync_dist=True) self.log("train_loss_lm", float(blip2_loss.loss_lm), batch_size=batch_size, sync_dist=True) self.log("train_loss", float(blip2_loss.loss), batch_size=batch_size, sync_dist=True) self.log("lr", self.trainer.optimizers[0].param_groups[0]['lr'], batch_size=batch_size, sync_dist=True) return blip2_loss.loss @staticmethod def add_model_specific_args(parent_parser): parser = parent_parser.add_argument_group("Stage1") # train mode parser.add_argument('--temperature', type=float, default=0.1, help='the temperature of NT_XentLoss') parser.add_argument('--save_every_n_epochs', type=int, default=0) parser.add_argument('--ptm', action='store_true', help='use graph-text matching or not', default=True) parser.add_argument('--lm', action='store_true', help='use language modeling or not', default=True) # evaluation parser.add_argument('--rerank_cand_num', type=int, default=128) # plm parser.add_argument('--plm_name', type=str, default='facebook/esm2_t30_150M_UR50D,可以是模型地址') parser.add_argument('--plm_tune', type=str, default='freeze') parser.add_argument('--load_4bit', action='store_true', default=False) parser.add_argument('--pool_size', type=int, default=0) # Bert parser.add_argument('--bert_hidden_dim', type=int, default=768, help='') parser.add_argument('--bert_name', type=str, default='microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract') parser.add_argument('--projection_dim', type=int, default=256) parser.add_argument('--cross_attention_freq', type=int, default=2) parser.add_argument('--num_query_token', type=int, default=8) # optimization parser.add_argument('--weight_decay', type=float, default=0.05, help='optimizer weight decay') parser.add_argument('--init_lr', type=float, default=1e-4, help='optimizer init learning rate') parser.add_argument('--min_lr', type=float, default=1e-5, help='optimizer min learning rate') parser.add_argument('--warmup_lr', type=float, default=1e-6, help='optimizer warmup learning rate') parser.add_argument('--warmup_steps', type=int, default=1000, help='optimizer warmup steps') parser.add_argument('--lr_decay_rate', type=float, default=0.9, help='optimizer lr decay rate') parser.add_argument('--scheduler', type=str, default='linear_warmup_cosine_lr', help='type of scheduler') # or linear_warmup_step_lr parser.add_argument('--init_checkpoint', type=str, default='') parser.add_argument('--retrieval_eval_epoch', type=int, default=10) return parent_parser @torch.no_grad() def eval_retrieval_fullset(prot_feat, text_feat, device): N = prot_feat.shape[0] B = 8 text_feat = text_feat.to(device) sim_p2t = [] for i in tqdm(range(0, N, B)): l_prot_feat = prot_feat[i:i+B].to(device) l_sim_q2t = (l_prot_feat.unsqueeze(1) @ text_feat.unsqueeze(-1)).squeeze() # shape = [B, 1, num_qs, D]; shape = [N, D, 1]; output shape = [B, N, num_qs] l_sim_p2t, _ = l_sim_q2t.max(-1) # shape = [B, N] sim_p2t.append(l_sim_p2t) sim_p2t = torch.cat(sim_p2t, dim=0).cpu() # shape = [N, N] rank_p2t = [] for i in range(0, N, B): sorted_ids = torch.argsort(sim_p2t[i:i+B].to(device), descending=True) rank_p2t.append((sorted_ids == torch.arange(i,i+sorted_ids.shape[0], device=device).reshape(-1, 1)).int().argmax(dim=-1)) rank_p2t = torch.cat(rank_p2t, dim=0) rank_t2p = [] for i in range(0, N, B): sorted_ids = torch.argsort(sim_p2t.T[i:i+B].to(device), descending=True) rank_t2p.append((sorted_ids == torch.arange(i,i+sorted_ids.shape[0], device=device).reshape(-1, 1)).int().argmax(dim=-1)) rank_t2p = torch.cat(rank_t2p, dim=0) p2t_acc = float((rank_p2t == 0).float().mean()) p2t_rec20 = float((rank_p2t < 20).float().mean()) t2p_acc = float((rank_t2p == 0).float().mean()) t2p_rec20 = float((rank_t2p < 20).float().mean()) p2t_acc = round(p2t_acc * 100, 2) p2t_rec20 = round(p2t_rec20 * 100, 2) t2p_acc = round(t2p_acc * 100, 2) t2p_rec20 = round(t2p_rec20 * 100, 2) return p2t_acc, p2t_rec20, t2p_acc, t2p_rec20, sim_p2t @torch.no_grad() def eval_retrieval_fullset_for_rerank(model, sim_p2t_total, prot_embed_total, prot_mask_total, text_total, text_mask_total, rerank_cand_num, device): N = sim_p2t_total.shape[0] B = 4 rcn = rerank_cand_num ## re-rank candidate numbers # print(f"sim_p2t_total shape: {sim_p2t_total.shape}") # print(f"rerank_cand_num: {rerank_cand_num}") rcn = min(rerank_cand_num, sim_p2t_total.shape[1]) hit_p2t = [] for i in tqdm(range(0, N, B), desc='re-ranking p2t'): sim = sim_p2t_total[i:i+B].to(device) rB = sim.shape[0] # real batch size topk_sim, topk_idx = sim.topk(k=rcn, dim=1) # shape = [B, rcn] topk_idx = topk_idx.cpu() prot_embed = prot_embed_total[i:i+B].to(device).repeat_interleave(rcn, 0) # shape = [B * rcn, num_qs, D] prot_mask = prot_mask_total[i:i+B].to(device).repeat_interleave(rcn, 0) # shape = [B * rcn, num_qs, D] text = text_total[topk_idx].flatten(0,1).to(device) # shape = [B * rcn, text_len] text_mask = text_mask_total[topk_idx].flatten(0,1).to(device) # shape = [B * rcn, text_len] ptm_sim = model.compute_ptm(prot_embed, prot_mask, text, text_mask).reshape(rB, rcn) ## fixme, using the linear clf's logits directly, without softmax sorted_ids = torch.argsort(topk_sim + ptm_sim, descending=True).cpu() # shape = [B, rcn] # sorted_ids = torch.argsort(gtm_sim, descending=True).cpu() # shape = [B, rcn] sorted_ids = torch.gather(topk_idx, 1, sorted_ids) # mapping to original ids hit_p2t.append((sorted_ids == torch.arange(i,i+rB).reshape(-1, 1)).int()) hit_p2t = torch.cat(hit_p2t, dim=0) # shape = [N, rcn] # p2t_acc = float((hit_p2t[:, 0]).float().mean()) # p2t_rec20 = float((hit_p2t[:, :20]).float().sum() / N) # print(p2t_acc, p2t_rec20) hit_t2p = [] sim_t2p_total = sim_p2t_total.T for i in tqdm(range(0, N, B), desc='re-ranking t2p'): sim = sim_t2p_total[i:i+B].to(device) rB = sim.shape[0] topk_sim, topk_idx = sim.topk(k=rcn, dim=1) topk_idx = topk_idx.cpu() text = text_total[i:i+B].to(device).repeat_interleave(rcn, 0) text_mask = text_mask_total[i:i+B].to(device).repeat_interleave(rcn, 0) prot_embed = prot_embed_total[topk_idx].to(device).flatten(0,1) prot_mask = prot_mask_total[topk_idx].to(device).flatten(0,1) ptm_sim = model.compute_ptm(prot_embed, prot_mask, text, text_mask).reshape(rB, rcn) sorted_ids = torch.argsort(topk_sim + ptm_sim, descending=True).cpu() # shape = [B, rcn] sorted_ids = torch.gather(topk_idx, 1, sorted_ids) hit_t2p.append((sorted_ids == torch.arange(i,i+sorted_ids.shape[0]).reshape(-1, 1)).int()) hit_t2p = torch.cat(hit_t2p, dim=0) p2t_acc = float((hit_p2t[:, 0]).float().mean()) p2t_rec20 = float((hit_p2t[:, :20]).float().sum() / N) t2p_acc = float((hit_t2p[:, 0]).float().mean()) t2p_rec20 = float((hit_t2p[:, :20]).float().sum() / N) p2t_acc = round(p2t_acc * 100, 2) p2t_rec20 = round(p2t_rec20 * 100, 2) t2p_acc = round(t2p_acc * 100, 2) t2p_rec20 = round(t2p_rec20 * 100, 2) return p2t_acc, p2t_rec20, t2p_acc, t2p_rec20 @torch.no_grad() def eval_retrieval_inbatch_with_rerank(model, dataloader, device=None): ''' include rerank ''' assert isinstance(model, Blip2Qformer) pad_token_id = model.tokenizer.pad_token_id model.eval() p2t_acc = 0 t2p_acc = 0 p2t_rec20 = 0 t2p_rec20 = 0 allcnt = 0 p2t_rerank_acc = 0 t2p_rerank_acc = 0 p2t_rerank_rec20 = 0 t2p_rerank_rec20 = 0 prot_feat_total = [] text_feat_total = [] prot_embed_total = [] prot_mask_total = [] text_total = [] text_mask_total = [] for batch in tqdm(dataloader): prot_batch, text_batch = batch prot_batch, text_batch = prot_batch.to(device), text_batch.to(device) text_total.append(text_batch.input_ids) text_mask_total.append(text_batch.attention_mask) prot_feats, prot_embeds = model.prot_forward(prot_batch) # shape = [B, num_qs, D] text_feats = model.text_forward(text_batch) # shape = [B, D] sim_q2t = (prot_feats.unsqueeze(1) @ text_feats.unsqueeze(-1)).squeeze() # shape = [B, 1, num_qs, D]; shape = [B, D, 1]; output shape = [B, B, num_qs] sim_p2t, _ = sim_q2t.max(-1) # shape = [B, B] B = sim_p2t.shape[0] sorted_ids = sim_p2t.argsort(descending=True).cpu() p2t_rank = (sorted_ids == torch.arange(B).reshape(-1, 1)).int().argmax(dim=-1) sorted_ids = sim_p2t.T.argsort(descending=True).cpu() t2p_rank = (sorted_ids == torch.arange(B).reshape(-1, 1)).int().argmax(dim=-1) p2t_acc += float((p2t_rank == 0).sum()) t2p_acc += float((t2p_rank == 0).sum()) p2t_rec20 += float((p2t_rank < 20).sum()) t2p_rec20 += float((t2p_rank < 20).sum()) allcnt += B prot_feat_total.append(prot_feats.cpu()) text_feat_total.append(text_feats.cpu()) prot_embed_total.append(prot_embeds.cpu()) prot_mask_total.append(prot_batch.attention_mask.cpu()) ## reranking prot_embeds = prot_embeds.repeat_interleave(B, 0) # shape = [B * B, prot_len, D] prot_mask = prot_batch.attention_mask.repeat_interleave(B, 0) # shape = [B * B, prot_len] text_ids = text_batch.input_ids.repeat(B, 1) # shape = [B * B, text_len] text_mask = text_batch.attention_mask.repeat(B, 1) # shape = [B * B, text_len] ## batched reranking batch_size = 64 ptm_sim = [] for i in range(0, prot_embeds.shape[0], batch_size): ptm_sim_local = model.compute_ptm(prot_embeds[i:i+batch_size], prot_mask[i:i+batch_size], text_ids[i:i+batch_size], text_mask[i:i+batch_size]) ptm_sim.append(ptm_sim_local) ptm_sim = torch.cat(ptm_sim, dim=0).reshape(B, B) rerank_sim = sim_p2t + ptm_sim ## p2t rerank sorted_ids = torch.argsort(rerank_sim, descending=True).cpu() # shape = [B, B] hit_p2t = (sorted_ids == torch.arange(B).reshape(-1, 1)).float() p2t_rerank_acc += float(hit_p2t[:, 0].sum()) p2t_rerank_rec20 += float(hit_p2t[:, :20].sum()) ## t2p rerank sorted_ids = torch.argsort(rerank_sim.T, descending=True).cpu() # shape = [B, B] hit_t2p = (sorted_ids == torch.arange(B).reshape(-1, 1)).float() t2p_rerank_acc += float(hit_t2p[:, 0].sum()) t2p_rerank_rec20 += float(hit_t2p[:, :20].sum()) prot_feat_total = torch.cat(prot_feat_total, dim=0) text_feat_total = torch.cat(text_feat_total, dim=0) prot_embed_total = pad_and_concat(prot_embed_total) prot_mask_total = pad_and_concat(prot_mask_total) text_total = pad_and_concat(text_total, fill_value=pad_token_id) text_mask_total = pad_and_concat(text_mask_total) # # text_total = torch.cat(text_total, dim=0) # text_mask_total = torch.cat(text_mask_total, dim=0) p2t_acc = round(p2t_acc/allcnt * 100, 2) t2p_acc = round(t2p_acc/allcnt * 100, 2) p2t_rec20 = round(p2t_rec20 / allcnt * 100, 2) t2p_rec20 = round(t2p_rec20 / allcnt * 100, 2) p2t_rerank_acc = round(p2t_rerank_acc / allcnt * 100, 2) t2p_rerank_acc = round(t2p_rerank_acc / allcnt * 100, 2) p2t_rerank_rec20 = round(p2t_rerank_rec20 / allcnt * 100, 2) t2p_rerank_rec20 = round(t2p_rerank_rec20 / allcnt * 100, 2) return p2t_acc, t2p_acc, p2t_rec20, t2p_rec20, \ p2t_rerank_acc, t2p_rerank_acc, p2t_rerank_rec20, t2p_rerank_rec20, \ prot_feat_total, text_feat_total, prot_embed_total, prot_mask_total, text_total, text_mask_total