Update handler.py

05261c9 verified 6 months ago

2.73 kB

from huggingface_inference_toolkit.logging import logger
import torch
from transformers import AutoTokenizer, AutoModelForMaskedLM

class EndpointHandler:
    def __init__(self, path=""):
        logger.info("Initializing Nucleotide Transformer...")
        self.model_id = "InstaDeepAI/nucleotide-transformer-v2-50m-multi-species"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
        self.model = AutoModelForMaskedLM.from_pretrained(self.model_id, trust_remote_code=True)
        
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.model.to(self.device)
        self.model.eval()
        logger.info(f"Model loaded on {self.device}")

    def _get_embedding(self, sequence):
        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)
            # mean of last hidden state -> [1, 512]
            return outputs.hidden_states[-1].mean(dim=1)

    def map_sensitivity(self, sequence, window_size, step):
        logger.info(f"Starting sensitivity map: window={window_size}, step={step}")
        baseline = self._get_embedding(sequence)
        
        results = []
        # Optimization: We can do this in a single list comprehension to keep it clean
        for i in range(0, len(sequence) - window_size, step or 1):
            # Create a 'mutated' version by reversing the window segment
            # (Reversing is more deterministic and faster than random.shuffle for testing)
            window = sequence[i : i + window_size]
            mutated_seq = sequence[:i] + window[::-1] + sequence[i + window_size:]
            
            mutated_emb = self._get_embedding(mutated_seq)
            
            # Distance calculation
            dist = torch.norm(baseline - mutated_emb).item()
            results.append({"coord": i, "score": dist})
            
        return results

    def __call__(self, data):
        """
        The Toolkit calls this method. 
        'data' is the dictionary from your payload.
        """
        logger.info(f"Payload received: {data}")
        
        inputs = data.get("inputs", "")
        method = data.get("method", "embed")
        # Ensure these are integers
        window_size = int(data.get("window_size", 50))
        step = int(data.get("step", 20)) # Default to 20 to match your test

        if method == "sensitivity":
            return self.map_sensitivity(inputs, window_size, step)
        
        # Default embedding behavior
        emb = self._get_embedding(inputs)
        return emb.cpu().tolist()[0]

	from huggingface_inference_toolkit.logging import logger
	import torch
	from transformers import AutoTokenizer, AutoModelForMaskedLM

	class EndpointHandler:
	def __init__(self, path=""):
	logger.info("Initializing Nucleotide Transformer...")
	self.model_id = "InstaDeepAI/nucleotide-transformer-v2-50m-multi-species"
	self.tokenizer = AutoTokenizer.from_pretrained(self.model_id, trust_remote_code=True)
	self.model = AutoModelForMaskedLM.from_pretrained(self.model_id, trust_remote_code=True)

	self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	self.model.to(self.device)
	self.model.eval()
	logger.info(f"Model loaded on {self.device}")

	def _get_embedding(self, sequence):
	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)
	# mean of last hidden state -> [1, 512]
	return outputs.hidden_states[-1].mean(dim=1)

	def map_sensitivity(self, sequence, window_size, step):
	logger.info(f"Starting sensitivity map: window={window_size}, step={step}")
	baseline = self._get_embedding(sequence)

	results = []
	# Optimization: We can do this in a single list comprehension to keep it clean
	for i in range(0, len(sequence) - window_size, step or 1):
	# Create a 'mutated' version by reversing the window segment
	# (Reversing is more deterministic and faster than random.shuffle for testing)
	window = sequence[i : i + window_size]
	mutated_seq = sequence[:i] + window[::-1] + sequence[i + window_size:]

	mutated_emb = self._get_embedding(mutated_seq)

	# Distance calculation
	dist = torch.norm(baseline - mutated_emb).item()
	results.append({"coord": i, "score": dist})

	return results

	def __call__(self, data):
	"""
	The Toolkit calls this method.
	'data' is the dictionary from your payload.
	"""
	logger.info(f"Payload received: {data}")

	inputs = data.get("inputs", "")
	method = data.get("method", "embed")
	# Ensure these are integers
	window_size = int(data.get("window_size", 50))
	step = int(data.get("step", 20)) # Default to 20 to match your test

	if method == "sensitivity":
	return self.map_sensitivity(inputs, window_size, step)

	# Default embedding behavior
	emb = self._get_embedding(inputs)
	return emb.cpu().tolist()[0]