|
|
|
|
|
from typing import Dict, Any |
|
|
|
from peft import PeftModel |
|
from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig |
|
import transformers |
|
|
|
|
|
import torch |
|
from torch import Tensor |
|
import torch.nn.functional as F |
|
from torch import nn |
|
|
|
from transformers.models.mistral.modeling_mistral import MistralAttention |
|
from ExtractableMistralAttention import forward |
|
|
|
MistralAttention.forward = forward |
|
|
|
class EndpointHandler(): |
|
def __init__(self, model_dir=''): |
|
self.instruction = 'Given a web search query, retrieve relevant passages that answer the query:\n' |
|
self.max_length = 4096 |
|
self.device = "cuda:0" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
self.tokenizer = AutoTokenizer.from_pretrained(model_dir, trust_remote_code=True) |
|
self.tokenizer.pad_token = '[PAD]' |
|
self.tokenizer.padding_side = 'left' |
|
|
|
bnb_config = BitsAndBytesConfig(load_in_8bit=True, bnb_8bit_compute_dtype=torch.float16) |
|
|
|
self.model = AutoModel.from_pretrained( |
|
model_dir, |
|
quantization_config=bnb_config, |
|
device_map="auto", |
|
trust_remote_code=True, |
|
attn_implementation="eager", |
|
) |
|
|
|
self.model.eval() |
|
|
|
|
|
def last_token_pool(self, last_hidden_states: Tensor, |
|
attention_mask: Tensor) -> Tensor: |
|
left_padding = (attention_mask[:, -1].sum() == attention_mask.shape[0]) |
|
if left_padding: |
|
return last_hidden_states[:, -1] |
|
else: |
|
sequence_lengths = attention_mask.sum(dim=1) - 1 |
|
batch_size = last_hidden_states.shape[0] |
|
return last_hidden_states[torch.arange(batch_size, device=last_hidden_states.device), sequence_lengths] |
|
|
|
|
|
def tokenize(self, text, request_type): |
|
if request_type == 'query': |
|
text = self.instruction + text |
|
return self.tokenizer(text + self.tokenizer.eos_token, max_length=self.max_length, truncation=True, return_tensors='pt').to(self.device) |
|
|
|
|
|
def extract_attn_vec(self, model): |
|
return self.model._modules['layers'][-1].self_attn.attn_vec |
|
|
|
|
|
def embed(self, text, request_type): |
|
tokens = self.tokenize(text, request_type) |
|
with torch.no_grad(): |
|
output = self.model(tokens['input_ids'], tokens['attention_mask']).last_hidden_state.detach() |
|
embedding = self.last_token_pool(output, tokens['attention_mask']) |
|
embedding = F.normalize(embedding, p=2, dim=1) |
|
|
|
attn_vec = self.extract_attn_vec(self.model) |
|
attn_vec = self.last_token_pool(attn_vec, tokens['attention_mask']) |
|
attn_vec = F.normalize(attn_vec, p=2, dim=1) |
|
del output, tokens |
|
torch.cuda.empty_cache() |
|
return embedding, attn_vec |
|
|
|
|
|
def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: |
|
""" |
|
data args: |
|
inputs (:obj: `str` | `PIL.Image` | `np.array`) |
|
kwargs |
|
Return: |
|
A :obj:`list` | `dict`: will be serialized and returned |
|
""" |
|
inputs = data.pop("inputs", data) |
|
id = inputs.pop("id", inputs) |
|
text = inputs.pop("text", inputs) |
|
request_type = inputs.pop("type", inputs) |
|
|
|
|
|
embeddings, attn_vec = self.embed(text, request_type) |
|
embeddings = embeddings[0].tolist() |
|
attn_vec = attn_vec[0].tolist() |
|
|
|
if request_type == 'query': |
|
return {"id": id, "embedding": embeddings, "attention_vec": attn_vec} |
|
|
|
elif request_type == 'document': |
|
return {"id": id, "embedding": embeddings} |
