| import torch |
| import torch.nn.functional as F |
|
|
| def warmup_lr(step): |
| return min(step, 5000) / 5000 |
|
|
|
|
| def training_with_esm(x_0_con, x_0_dis, trainer, trainer_dis, ns_con, ns_dis, trans, FLAGS, cls_head, class_weights): |
|
|
| t = torch.randint(FLAGS.T, size=(x_0_con['input_ids'].shape[0], ), device=x_0_con['input_ids'].device) |
| pt = torch.ones_like(t).float() / FLAGS.T |
|
|
| embeddings = trainer.model.get_embeddings(x_0_con) |
|
|
| noise = torch.randn_like(embeddings) |
| x_t_con = trainer.make_x_t(embeddings, t, noise) |
| torch.cuda.empty_cache() |
| log_x_start = torch.log(x_0_dis.float().clamp(min=1e-30)) |
| x_t_dis = trainer_dis.q_sample(log_x_start=log_x_start, t=t) |
|
|
| eps = trainer.model(x_t_con, t, x_t_dis.to(x_t_con.device)) |
| ps_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps) |
| con_loss = F.mse_loss(eps, noise, reduction='none') |
| con_loss = con_loss.mean() |
| kl, ps_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, x_t_con) |
| ps_0_dis = torch.exp(ps_0_dis) |
| kl_prior = trainer_dis.kl_prior(log_x_start) |
| dis_loss = (kl / pt + kl_prior).mean() |
| torch.cuda.empty_cache() |
| noise_ns = torch.randn_like(ns_con) |
| ns_t_con = trainer.make_x_t(ns_con, t, noise_ns) |
| log_ns_start = torch.log(ns_dis.float().clamp(min=1e-30)) |
| ns_t_dis = trainer_dis.q_sample(log_x_start=log_ns_start, t=t) |
| torch.cuda.empty_cache() |
| eps_ns = trainer.model(x_t_con, t, ns_t_dis.to(ns_t_dis.device)) |
| ns_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps_ns) |
| _, ns_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, ns_t_con) |
| ns_0_dis = torch.exp(ns_0_dis) |
|
|
|
|
| torch.cuda.empty_cache() |
| triplet_loss = torch.nn.TripletMarginLoss(margin=0.5, p=2) |
| triplet_con = triplet_loss(embeddings, ps_0_con, ns_0_con) |
| total_loss = con_loss |
| st=0 |
| triplet_dis = [] |
| for item in trans.output_info: |
| ed = st + item[0] |
| ps_dis = F.cross_entropy(ps_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none') |
| ns_dis = F.cross_entropy(ns_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none') |
|
|
| triplet_dis.append(max((ps_dis-ns_dis).mean()+1,0)) |
| st = ed |
| triplet_dis = sum(triplet_dis)/len(triplet_dis) |
|
|
|
|
| return total_loss, triplet_con, dis_loss, triplet_dis |
|
|
|
|
| def make_negative_condition_esm(x_0_con, x_0_dis, base_model): |
| device = x_0_dis.device |
|
|
| with torch.no_grad(): |
| _, _, embeddings, _ = base_model( |
| input_ids=x_0_con['input_ids'], |
| attention_masks=x_0_con['attention_mask'] |
| ) |
|
|
| B = embeddings.shape[0] |
|
|
| indices = torch.randperm(B) |
|
|
| ns_con = embeddings[indices].clone() |
| ns_dis = x_0_dis[indices].clone() |
|
|
| return ns_con, ns_dis |
|
|
