|
from diffusers import DiffusionPipeline |
|
import torch |
|
from sklearn.neighbors import KNeighborsRegressor |
|
import numpy as np |
|
|
|
class inDelphiPipeline(DiffusionPipeline): |
|
def __init__(self, inDelphi_model, onebp_features, insert_probabilities, m654): |
|
super().__init__() |
|
|
|
self.register_modules(inDelphi_model=inDelphi_model) |
|
self.onebp_feature_mean = onebp_features.mean(axis=0) |
|
self.onebp_feature_std = onebp_features.std(axis=0) |
|
self.insertion_model = KNeighborsRegressor(weights='distance').fit((onebp_features - self.onebp_feature_mean) / self.onebp_feature_std, insert_probabilities) |
|
self.m654 = m654 / np.maximum(np.linalg.norm(m654, ord=1, axis=1, keepdims=True), 1e-6) |
|
self.m4 = m654.reshape(16, 4, 4).sum(axis=0) |
|
self.m4 = self.m4 / np.maximum(np.linalg.norm(self.m4, ord=1, axis=1, keepdims=True), 1e-6) |
|
|
|
@torch.no_grad() |
|
def __call__(self, batch, use_m654=False): |
|
mh_weights, mhless_weights, total_del_len_weights = self.inDelphi_model( |
|
batch["mh_input"].to(self.inDelphi_model.device), |
|
batch["mh_del_len"].to(self.inDelphi_model.device) |
|
).values() |
|
mX = self.m654 if use_m654 else self.m4 |
|
log_total_weights = total_del_len_weights.sum(dim=1, keepdim=True).log() |
|
precisions = 1 - torch.distributions.Categorical(total_del_len_weights[:,:28]).entropy() / torch.log(torch.tensor(28)) |
|
onebp_features = torch.cat([ |
|
batch["onebp_feature"], |
|
precisions[:, None].cpu(), |
|
log_total_weights.cpu() |
|
], dim=1).cpu().numpy() |
|
pre_insert_probabilities = self.insertion_model.predict((onebp_features - self.onebp_feature_mean) / self.onebp_feature_std) |
|
pre_insert_1bps = mX[batch['m654'] % 4] if mX.shape[0] == 4 else mX[batch['m654']] |
|
return { |
|
"mh_weight": [ |
|
mh_weights[i, batch["mh_del_len"][i] < self.inDelphi_model.config.DELLEN_LIMIT] |
|
for i in range(len(batch["mh_del_len"])) |
|
], |
|
"mhless_weight": mhless_weights, |
|
"total_del_len_weight": total_del_len_weights, |
|
"pre_insert_probability": pre_insert_probabilities, |
|
"pre_insert_1bp": pre_insert_1bps |
|
} |
