Update handler.py

handler.py

@@ -1,47 +1,86 @@
-import os
 import torch
-from typing import Any, Dict, List
-from transformers import AutoConfig, AutoTokenizer, AutoModelForMaskedLM
-
-# Force the trust flag globally
-os.environ["HF_HUB_TRUST_REMOTE_CODE"] = "True"
+from transformers import AutoTokenizer, AutoModelForMaskedLM
+import numpy as np
 
 class EndpointHandler:
     def __init__(self, path=""):
-        # We ignore the local 'path' and pull fresh from the source
+        # Load model and tokenizer
         self.model_id = "InstaDeepAI/nucleotide-transformer-v2-50m-multi-species"
-        
-        self.config = AutoConfig.from_pretrained(self.model_id, trust_remote_code=True)
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
-        self.model = AutoModelForMaskedLM.from_pretrained(
-            self.model_id, 
-            config=self.config, 
-            trust_remote_code=True
-        )
+        self.model = AutoModelForMaskedLM.from_pretrained(self.model_id, trust_remote_code=True)
         
-        if torch.cuda.is_available():
-            self.model = self.model.to("cuda")
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
         self.model.eval()
 
-    def __call__(self, data: Dict[str, Any]) -> List[float]:
-        inputs = data.pop("inputs", data)
-        if isinstance(inputs, list):
-            inputs = inputs[0]
+    def _get_embedding(self, sequence):
+        """Helper to get a single mean embedding."""
+        inputs = self.tokenizer(sequence, return_tensors="pt", truncation=True, max_length=1024).to(self.device)
+        with torch.no_grad():
+            outputs = self.model(**inputs, output_hidden_states=True)
+            # Use mean of the last hidden state
+            embeddings = outputs.hidden_states[-1].mean(dim=1)
+        return embeddings
 
-        # 12.2kb APRIL promoter chunking
-        chunk_size = 1000
-        stride = 500
-        chunks = [inputs[i:i + chunk_size] for i in range(0, len(inputs), stride)]
+    def map_sensitivity(self, sequence, window_size=50, step=100):
+        """Generates a sensitivity map by perturbing segments of the sequence."""
+        # 1. Get Baseline
+        baseline_embedding = self._get_embedding(sequence)
         
-        all_embeddings = []
-        with torch.no_grad():
-            for chunk in chunks:
-                tokens = self.tokenizer(chunk, return_tensors='pt', padding=True, truncation=True, max_length=chunk_size)
-                if torch.cuda.is_available():
-                    tokens = {k: v.to("cuda") for k, v in tokens.items()}
+        # 2. Create variants
+        variants = []
+        indices = []
+        seq_list = list(sequence)
+        
+        for i in range(0, len(sequence) - window_size, step):
+            # Create a "shuffled" variant of the window
+            variant_seq = seq_list.copy()
+            sub_seq = variant_seq[i : i + window_size]
+            import random
+            random.shuffle(sub_seq)
+            variant_seq[i : i + window_size] = sub_seq
+            
+            variants.append("".join(variant_seq))
+            indices.append(i)
+
+        # 3. Batch Inference (Processing variants in chunks to fit in VRAM)
+        batch_size = 16 
+        distances = []
+        
+        for k in range(0, len(variants), batch_size):
+            batch_texts = variants[k : k + batch_size]
+            inputs = self.tokenizer(batch_texts, return_tensors="pt", padding=True, truncation=True, max_length=1024).to(self.device)
+            
+            with torch.no_grad():
+                outputs = self.model(**inputs, output_hidden_states=True)
+                batch_embeddings = outputs.hidden_states[-1].mean(dim=1)
                 
-                outputs = self.model(**tokens, output_hidden_states=True)
-                chunk_emb = torch.mean(outputs.hidden_states[-1], dim=1).squeeze()
-                all_embeddings.append(chunk_emb)
+                # Calculate Euclidean Distance to baseline on GPU
+                # distance = sqrt(sum((a - b)^2))
+                diff = batch_embeddings - baseline_embedding
+                dist = torch.norm(diff, dim=1)
+                distances.extend(dist.cpu().tolist())
+
+        # 4. Return coordinates and their corresponding sensitivity scores
+        return [{"coord": idx, "score": score} for idx, score in zip(indices, distances)]
+
+    def __call__(self, data):
+        """
+        Args:
+            data (:obj:`dict`):
+                - "inputs": the DNA sequence
+                - "method": "embed" (default) or "sensitivity"
+        """
+        inputs = data.get("inputs", "")
+        method = data.get("method", "embed")
+
+        if not inputs:
+            return {"error": "No input sequence provided"}
 
-        return torch.stack(all_embeddings).mean(dim=0).cpu().numpy().tolist()
+        if method == "sensitivity":
+            # Returns the map of high-leverage coordinates
+            return self.map_sensitivity(inputs)
+        else:
+            # Standard embedding behavior
+            embedding = self._get_embedding(inputs)
+            return embedding.cpu().tolist()[0]