Upload pipeline.py with huggingface_hub

pipeline.py

@@ -0,0 +1,30 @@
+import torch
+from diffusers import DiffusionPipeline
+import torch.nn.functional as F
+import numpy as np
+
+class FOREcasTPipeline(DiffusionPipeline):
+    def __init__(self, FOREcasT_model, MAX_DEL_SIZE):
+        super().__init__()
+
+        self.register_modules(FOREcasT_model=FOREcasT_model)
+        self.MAX_DEL_SIZE = MAX_DEL_SIZE
+        self.lefts = np.concatenate([
+            np.arange(-DEL_SIZE, 1)
+            for DEL_SIZE in range(self.MAX_DEL_SIZE, -1, -1)
+        ] + [np.zeros(20, np.int64)])
+        self.rights = np.concatenate([
+            np.arange(0, DEL_SIZE + 1)
+            for DEL_SIZE in range(self.MAX_DEL_SIZE, -1, -1)
+        ] + [np.zeros(20, np.int64)])
+        self.inss = (self.MAX_DEL_SIZE + 2) * (self.MAX_DEL_SIZE + 1) // 2 * [""] + ["A", "C", "G", "T", "AA", "AC", "AG", "AT", "CA", "CC", "CG", "CT", "GA", "GC", "GG", "GT", "TA", "TC", "TG", "TT"]
+
+    @torch.no_grad()
+    def __call__(self, batch):
+        assert batch["feature"].shape[1] == len(self.lefts), "the possible mutation number of the input feature does not fit the pipeline"
+        return {
+            "proba": F.softmax(self.FOREcasT_model(batch["feature"].to(self.FOREcasT_model.device))["logit"], dim=-1),
+            "left": self.lefts,
+            "right": self.rights,
+            "ins_seq": self.inss
+        }