upload

dpr_biencoder.38.602

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

dump.py

@@ -0,0 +1,114 @@
+import collections
+
+from torch.serialization import default_restore_location
+
+from transformers import BertTokenizer, BertModel
+import torch
+from torch.utils.data import Dataset, DataLoader
+from tqdm import tqdm
+import numpy as np
+import pickle
+import argparse
+import csv
+
+nq_temp = {}
+
+CheckpointState = collections.namedtuple("CheckpointState",
+                                         ['model_dict', 'optimizer_dict', 'scheduler_dict', 'offset', 'epoch',
+                                          'encoder_params'])
+
+def load_states_from_checkpoint(model_file: str) -> CheckpointState:
+    state_dict = torch.load(model_file, map_location=lambda s, l: default_restore_location(s, 'cpu'))
+    return CheckpointState(**state_dict)
+
+class DocPool(Dataset):
+    def __init__(self, path):
+        doc = []
+        with open(path, "r", encoding="utf8") as f:
+            lines = csv.reader(f, delimiter='\t')
+            for _id, _text in lines:
+                doc.append(_text)
+        self.doc = doc
+        
+
+    def __len__(self):
+        return len(self.doc)
+
+    def __getitem__(self, index):
+        doc = self.doc[index]
+        return index, doc
+
+
+def my_collate(batch):
+    batch = list(zip(*batch))
+    res = {'id': batch[0], 'doc': batch[1]}
+    del batch
+    return res
+
+
+def extract_feature(args):
+    torch.manual_seed(2024)
+    torch.cuda.manual_seed(2024)
+    np.random.seed(2024)
+    if(args.doc_or_query == 'doc'):
+        _path = './nq_doc.tsv'
+        _out_path = './doc_embedding.pickle'
+        _prefix = 'ctx_model.'
+    else:
+        _path = './nq_query.tsv'
+        _out_path = './query_embedding.pickle'
+        _prefix = 'question_model.'
+    with torch.no_grad():
+        doc_dataset = DocPool(_path)
+        print(len(doc_dataset))
+        doc_dataloader = DataLoader(dataset=doc_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=my_collate)
+        tokenizer = BertTokenizer.from_pretrained(args.model_name)
+        model = BertModel.from_pretrained(args.model_name, return_dict=True)
+        saved_state = load_states_from_checkpoint(args.model_file)
+        prefix_len = len(_prefix)
+        ctx_state = {key[prefix_len:]: value for (key, value) in saved_state.model_dict.items() if
+                     key.startswith(_prefix)}
+        model.load_state_dict(ctx_state, strict=False)
+        model = torch.nn.DataParallel(model)
+        device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+        model = model.to(device)
+        model.eval()
+
+        ids = []
+        idx = 0
+        doc_feature = np.zeros((len(doc_dataset), 768))
+
+        for batch_data in tqdm(doc_dataloader):
+            doc_id = batch_data['id']
+            doc_body = batch_data['doc']
+            inputs = tokenizer(doc_body, padding=True, truncation=True, return_tensors="pt",
+                               add_special_tokens=True).to(device)
+            outputs = model(**inputs)
+            pooler_output = outputs.last_hidden_state[:, 0]
+
+            ids.extend(doc_id)
+            doc_feature[idx: idx + pooler_output.shape[0]] = pooler_output.cpu().numpy()
+            idx += pooler_output.shape[0]
+
+        feature_dic = {}
+        for i, id_i in enumerate(ids):
+            feature_dic[id_i] = doc_feature[i]
+        with open(_out_path, 'wb') as f:
+            pickle.dump(feature_dic, f)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--batch_size', type=int, default=512,
+                        help='minibatch size')
+    parser.add_argument('--model_name', type=str, default='Luyu/co-condenser-wiki',
+                        help='model name')
+    parser.add_argument('--model_file', type=str, default='./dpr_biencoder.38.602',
+                        help='model name')
+    parser.add_argument('--doc_or_query', type=str, default='query',
+                        help='transfer documents or queries')
+
+    args = parser.parse_args()
+    extract_feature(args)