Upload 2 files

handler.py

@@ -0,0 +1,99 @@
+# handler file for Huggingface Inference API
+
+from typing import Dict, Any
+
+from peft import PeftModel
+from transformers import AutoTokenizer, AutoModel, BitsAndBytesConfig
+import transformers
+from transformers.models.mistral.modeling_mistral import MistralAttention
+from ExtractableMistralAttention import forward
+
+MistralAttention.forward = forward
+
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+
+
+class EndpointHandler():
+    def __init__(self):
+        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('intfloat/e5-mistral-7b-instruct', 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(
+            '',
+            quantization_config=bnb_config,
+            device_map="auto",
+            trust_remote_code=True,
+            attn_implementation="eager",
+        )
+        self.model = PeftModel.from_pretrained(model, '/lora')
+        self.model.eval()
+
+
+    def last_token_pool(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, type):
+        if 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(model):
+        return model._modules['layers'][-1].self_attn.attn_vec
+
+
+    def embed(self, text, type):
+        tokens = self.tokenize(text, 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
+        """
+
+        id = data.pop("id", data)
+        text = data.pop("text", data)
+        type = data.pop("type", data)
+
+        embeddings, attn_vec = self.embed(text, type)
+        embeddings = embeddings[0].tolist()
+        attn_vec = attn_vec[0].tolist()
+
+        if type == 'query':
+            return {"id": id, "embedding": embeddings, "attention_vec": attn_vec}
+        
+        elif type == 'document':
+            return {"id": id, "embedding": embeddings}

requirements.txt

@@ -0,0 +1,2 @@
+peft
+torch