Update scripts/gte_embedding.py

scripts/gte_embedding.py

@@ -1,42 +1,25 @@
-import logging
-from typing import Dict, Optional, List, Tuple
-import os
-
-import heapq
-import json
-import logging
-import os
-import queue
-import sys
-import time
-from tqdm import tqdm
+# coding=utf-8
+# Copyright 2024 The GTE Team Authors and Alibaba Group.
+# Licensed under the Apache License, Version 2.0 (the "License");
 
-import torch
 from collections import defaultdict
-from torch.utils.data._utils.worker import ManagerWatchdog
-import numpy as np
-import torch.distributed as dist
-from torch import nn, Tensor
-import torch.nn.functional as F
-from transformers import AutoModel, AutoTokenizer
-from transformers.file_utils import ModelOutput
+from typing import Dict, List, Tuple
 
-logger = logging.getLogger(__name__)
+import numpy as np
+import torch
+from transformers import AutoModelForTokenClassification, AutoTokenizer
+from transformers.utils import is_torch_npu_available
 
 
-class GTEEmbeddidng(nn.Module):
+class GTEEmbeddidng(torch.nn.Module):
     def __init__(self,
                  model_name: str = None,
                  normalized: bool = True,
-                 pooling_method: str = 'cls',
                  use_fp16: bool = True,
                  device: str = None
                 ):
         super().__init__()
-        self.load_model(model_name)
-        self.vocab_size = self.model.config.vocab_size
         self.normalized = normalized
-        self.pooling_method = pooling_method
         if device:
             self.device = torch.device(device)
         else:
@@ -49,40 +32,13 @@ class GTEEmbeddidng(nn.Module):
             else:
                 self.device = torch.device("cpu")
                 use_fp16 = False
-        self.model.to(self.device)
-        self.sparse_linear.to(self.device)
-        if use_fp16:
-            self.model.half()
-            self.sparse_linear.half()
-
-    def load_model(self, model_name):
-        if not os.path.exists(model_name):
-            cache_folder = os.getenv('HF_HUB_CACHE')
-            model_name = snapshot_download(repo_id=model_name,
-                                           cache_dir=cache_folder,
-                                           ignore_patterns=['flax_model.msgpack', 'rust_model.ot', 'tf_model.h5'])
-
-        self.model = AutoModel.from_pretrained(model_name, trust_remote_code=True)
-        self.sparse_linear = torch.nn.Linear(in_features=self.model.config.hidden_size, out_features=1)
+        self.use_fp16 = use_fp16
         self.tokenizer = AutoTokenizer.from_pretrained(model_name)
-        self.model.eval()
-        if os.path.exists(os.path.join(model_name, 'sparse_linear.pt')):
-            logger.info('loading existing sparse_linear---------')
-            self.load_pooler(model_dir=model_name)
-        else:
-            logger.warring('The parameters of  sparse linear is not found')
-
-    def dense_embedding(self, hidden_state, mask):
-        if self.pooling_method == 'cls':
-            return hidden_state[:, 0]
-        elif self.pooling_method == 'mean':
-            s = torch.sum(hidden_state * mask.unsqueeze(-1).float(), dim=1)
-            d = mask.sum(axis=1, keepdim=True).float()
-            return s / d
-
-    def sparse_embedding(self, hidden_state, input_ids, return_embedding: bool = True):
-        token_weights = torch.relu(self.sparse_linear(hidden_state))
-        return token_weights
+        self.model = AutoModelForTokenClassification.from_pretrained(
+            model_name, trust_remote_code=True, torch_dtype=torch.float16 if self.use_fp16 else None
+        )
+        self.vocab_size = self.model.config.vocab_size
+        self.model.to(self.device)
 
     def _process_token_weights(self, token_weights: np.ndarray, input_ids: list):
         # conver to dict
@@ -127,7 +83,7 @@ class GTEEmbeddidng(nn.Module):
 
     @torch.no_grad()
     def _encode(self,
-                texts: Dict[str, Tensor] = None,
+                texts: Dict[str, torch.Tensor] = None,
                 dimension: int = None,
                 max_length: int = 1024,
                 batch_size: int = 16,
@@ -136,27 +92,22 @@ class GTEEmbeddidng(nn.Module):
 
         text_input = self.tokenizer(texts, padding=True, truncation=True, return_tensors='pt', max_length=max_length)
         text_input = {k: v.to(self.model.device) for k,v in text_input.items()}
-        last_hidden_state = self.model(**text_input, return_dict=True).last_hidden_state
+        model_out = self.model(**text_input, return_dict=True)
 
         output = {}
         if return_dense:
-            dense_vecs = self.dense_embedding(last_hidden_state, text_input['attention_mask'])
-            dense_vecs = dense_vecs[:, :dimension]
+            dense_vecs = model_out.last_hidden_state[:, 0, :dimension]
             if self.normalized:
                 dense_vecs = torch.nn.functional.normalize(dense_vecs, dim=-1)
             output['dense_embeddings'] = dense_vecs
         if return_sparse:
-            token_weights = self.sparse_embedding(last_hidden_state, text_input['input_ids']).squeeze(-1)
+            token_weights = torch.relu(model_out.logits).squeeze(-1)
             token_weights = list(map(self._process_token_weights, token_weights.detach().cpu().numpy().tolist(),
                                                     text_input['input_ids'].cpu().numpy().tolist()))
             output['token_weights'] = token_weights
 
         return output
 
-    def load_pooler(self, model_dir):
-        sparse_state_dict = torch.load(os.path.join(model_dir, 'sparse_linear.pt'), map_location='cpu')
-        self.sparse_linear.load_state_dict(sparse_state_dict)
-    
     def _compute_sparse_scores(self, embs1, embs2):
         scores = 0
         for token, weight in embs1.items():
@@ -188,3 +139,16 @@ class GTEEmbeddidng(nn.Module):
             self.compute_sparse_scores(embs1['token_weights'], embs2['token_weights']) * sparse_weight
         scores = scores.tolist()
         return scores
+
+
+if __name__ == '__main__':
+    gte = GTEEmbeddidng('Alibaba-NLP/gte-multilingual-base')
+    docs =  [
+        "黑龙江离俄罗斯很近",
+        "哈尔滨是中国黑龙江省的省会，位于中国东北",
+        "you are the hero"
+    ]
+    print('docs', docs)
+    embs = gte.encode(docs, return_dense=True,return_sparse=True)
+    print('dense vecs', embs['dense_embeddings'])
+    print('sparse vecs', embs['token_weights'])