cr

splade_query/handler.py → handler.py

@@ -3,40 +3,54 @@ from transformers import AutoTokenizer, AutoModelForMaskedLM
 import torch
 from subprocess import run
 
+
 # set device
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# set path
+query_emb_model_path = "/splade_query"
+doc_emb_model_path = "/splade_doc"
+
 
 class EndpointHandler():
     def __init__(self, path=""):
-        # self.pipeline = pipeline("text-classification", model=path)
-        # self.holidays = holidays.US()
-        self.query_model = AutoModelForMaskedLM.from_pretrained(path).to(device)
-        self.query_tokenizer = AutoTokenizer.from_pretrained(path)
+        self.query_model = AutoModelForMaskedLM.from_pretrained(path+query_emb_model_path).to(device)
+        self.query_tokenizer = AutoTokenizer.from_pretrained(path+query_emb_model_path)
+        self.doc_model = AutoModelForMaskedLM.from_pretrained(path+doc_emb_model_path).to(device)
+        self.doc_tokenizer = AutoTokenizer.from_pretrained(path+doc_emb_model_path)
 
 
     def __call__(self, data: Dict[str, Any]) -> Tuple[List[List[int]], List[List[float]]]:
         """
-        data args:
-            inputs (:obj: `str`)
-            date (:obj: `str`)
-        Return:
+       data args:
+            inputs (:obj: `List[str]`)
+            task (:obj: `str`)
+      Return:
             A :obj:`list` | `dict`: will be serialized and returned
         """
         # get inputs
         texts = data.pop("inputs", data)
-
-        tokens = self.query_tokenizer(
+        task = data.pop("task", data)
+        emb_model = None
+        tokenizer = None
+        
+        if task == "query_emb":
+            emb_model = self.query_model
+            tokenizer = self.query_tokenizer
+        elif task == "doc_emb":
+            emb_model = self.doc_model
+            tokenizer = self.doc_tokenizer
+        else:
+            raise ValueError("task must be either 'query_emb' or 'doc_emb'")
+        
+        tokens = tokenizer(
             texts, truncation=True, padding=True, return_tensors="pt"
         )
 
-        tokens = self.query_tokenizer(
-            texts, truncation=True, padding=True, return_tensors="pt"
-        )
         if torch.cuda.is_available():
             tokens = tokens.to("cuda")
 
-        output = self.query_model(**tokens)
+        output = emb_model(**tokens)
         logits, attention_mask = output.logits, tokens.attention_mask
         relu_log = torch.log(1 + torch.relu(logits))
         weighted_log = relu_log * attention_mask.unsqueeze(-1)

splade_doc/README.md

@@ -0,0 +1,43 @@
+---
+license: cc-by-nc-sa-4.0
+language: "en"
+tags:
+- splade
+- query-expansion
+- document-expansion
+- bag-of-words
+- passage-retrieval
+- knowledge-distillation
+- document encoder
+datasets:
+- ms_marco
+---
+## Efficient SPLADE 
+Efficient SPLADE model for passage retrieval. This architecture uses two distinct models for query and document inference. This is the **doc** one, please also download the **query** one (https://huggingface.co/naver/efficient-splade-VI-BT-large-query). For additional details, please visit:
+* paper: https://dl.acm.org/doi/10.1145/3477495.3531833
+* code: https://github.com/naver/splade
+| | MRR@10 (MS MARCO dev) | R@1000 (MS MARCO dev) | Latency (PISA) ms | Latency (Inference) ms
+| --- | --- | --- | --- | --- |
+| `naver/efficient-splade-V-large` | 38.8 | 98.0 | 29.0 | 45.3
+| `naver/efficient-splade-VI-BT-large` | 38.0 | 97.8 | 31.1 | 0.7
+## Citation
+If you use our checkpoint, please cite our work:
+```
+@inproceedings{10.1145/3477495.3531833,
+author = {Lassance, Carlos and Clinchant, St\'{e}phane},
+title = {An Efficiency Study for SPLADE Models},
+year = {2022},
+isbn = {9781450387323},
+publisher = {Association for Computing Machinery},
+address = {New York, NY, USA},
+url = {https://doi.org/10.1145/3477495.3531833},
+doi = {10.1145/3477495.3531833},
+abstract = {Latency and efficiency issues are often overlooked when evaluating IR models based on Pretrained Language Models (PLMs) in reason of multiple hardware and software testing scenarios. Nevertheless, efficiency is an important part of such systems and should not be overlooked. In this paper, we focus on improving the efficiency of the SPLADE model since it has achieved state-of-the-art zero-shot performance and competitive results on TREC collections. SPLADE efficiency can be controlled via a regularization factor, but solely controlling this regularization has been shown to not be efficient enough. In order to reduce the latency gap between SPLADE and traditional retrieval systems, we propose several techniques including L1 regularization for queries, a separation of document/query encoders, a FLOPS-regularized middle-training, and the use of faster query encoders. Our benchmark demonstrates that we can drastically improve the efficiency of these models while increasing the performance metrics on in-domain data. To our knowledge, we propose the first neural models that, under the same computing constraints, achieve similar latency (less than 4ms difference) as traditional BM25, while having similar performance (less than 10% MRR@10 reduction) as the state-of-the-art single-stage neural rankers on in-domain data.},
+booktitle = {Proceedings of the 45th International ACM SIGIR Conference on Research and Development in Information Retrieval},
+pages = {2220–2226},
+numpages = {7},
+keywords = {splade, latency, information retrieval, sparse representations},
+location = {Madrid, Spain},
+series = {SIGIR '22}
+}
+```

splade_doc/config.json

@@ -0,0 +1,24 @@
+{
+  "_name_or_path": "/tmp-network/user/classanc/CoCodenser/flops_mlm_together/10_epochs/distilbert_256_64_0.001/",
+  "activation": "gelu",
+  "architectures": [
+    "DistilBertForMaskedLM"
+  ],
+  "attention_dropout": 0.1,
+  "dim": 768,
+  "dropout": 0.1,
+  "hidden_dim": 3072,
+  "initializer_range": 0.02,
+  "max_position_embeddings": 512,
+  "model_type": "distilbert",
+  "n_heads": 12,
+  "n_layers": 6,
+  "pad_token_id": 0,
+  "qa_dropout": 0.1,
+  "seq_classif_dropout": 0.2,
+  "sinusoidal_pos_embds": false,
+  "tie_weights_": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.11.3",
+  "vocab_size": 30522
+}

splade_doc/pytorch_model.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ee43cec6d21979da10d6520dbe1705a251debe690878ccebe2d6e9e9c3b0930e
+size 267982639

splade_doc/special_tokens_map.json

@@ -0,0 +1 @@
+{"unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]"}

splade_doc/tokenizer_config.json

@@ -0,0 +1 @@
+{"do_lower_case": true, "unk_token": "[UNK]", "sep_token": "[SEP]", "pad_token": "[PAD]", "cls_token": "[CLS]", "mask_token": "[MASK]", "tokenize_chinese_chars": true, "strip_accents": null, "do_basic_tokenize": true, "never_split": null, "model_max_length": 512, "special_tokens_map_file": null, "name_or_path": "/tmp-network/user/classanc/CoCodenser/flops_mlm_together/10_epochs/distilbert_256_64_0.001/", "tokenizer_class": "DistilBertTokenizer"}