Update handler.py

3b7c54f over 3 years ago

1.87 kB

from typing import Dict, List, Any
from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

class EndpointHandler:
    def __init__(self, path: str = ""):
        # Device to use for inference. "cuda" if CUDA is available, else "cpu"
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        
        # Load tokenizer and model
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        self.model = AutoModelForMaskedLM.from_pretrained(path).to(self.device)

    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
        """
        data args:
            inputs (:obj: `str`)
        Return:
            A :obj:`list` | `dict`: will be serialized and returned
        """
        # Get inputs
        inputs = data.pop("inputs", data)
        if isinstance(inputs, str):
            inputs = [inputs]

        # Run the model to get logits
        inputs = self.tokenizer(
            inputs,
            return_tensors="pt",
            padding=True,
            truncation=True,
            max_length=self.model.config.max_position_embeddings,  # Max length according to the model configuration
        ).to(self.device)
        with torch.no_grad():
            logits = self.model(**inputs).logits

        # Apply transformations to get the sparse vectors
        inter = torch.log1p(torch.relu(logits))
        token_max = torch.max(inter, dim=1)

        nz_tokens_i, nz_tokens_j = torch.where(token_max.values > 0)

        output = []
        for i in range(token_max.values.shape[0]):
            nz_tokens = nz_tokens_j[nz_tokens_i == i]
            nz_weights = token_max.values[i, nz_tokens]
            output.append(
                {"indices": nz_tokens.tolist(), "values": nz_weights.tolist()}
            )

        return output[0] if isinstance(inputs, str) else output

	from typing import Dict, List, Any
	from transformers import AutoTokenizer, AutoModelForMaskedLM
	import torch

	class EndpointHandler:
	def __init__(self, path: str = ""):
	# Device to use for inference. "cuda" if CUDA is available, else "cpu"
	self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	# Load tokenizer and model
	self.tokenizer = AutoTokenizer.from_pretrained(path)
	self.model = AutoModelForMaskedLM.from_pretrained(path).to(self.device)

	def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
	"""
	data args:
	inputs (:obj: `str`)
	Return:
	A :obj:`list` \| `dict`: will be serialized and returned
	"""
	# Get inputs
	inputs = data.pop("inputs", data)
	if isinstance(inputs, str):
	inputs = [inputs]

	# Run the model to get logits
	inputs = self.tokenizer(
	inputs,
	return_tensors="pt",
	padding=True,
	truncation=True,
	max_length=self.model.config.max_position_embeddings, # Max length according to the model configuration
	).to(self.device)
	with torch.no_grad():
	logits = self.model(**inputs).logits

	# Apply transformations to get the sparse vectors
	inter = torch.log1p(torch.relu(logits))
	token_max = torch.max(inter, dim=1)

	nz_tokens_i, nz_tokens_j = torch.where(token_max.values > 0)

	output = []
	for i in range(token_max.values.shape[0]):
	nz_tokens = nz_tokens_j[nz_tokens_i == i]
	nz_weights = token_max.values[i, nz_tokens]
	output.append(
	{"indices": nz_tokens.tolist(), "values": nz_weights.tolist()}
	)

	return output[0] if isinstance(inputs, str) else output