hannibal046 / nanocolbert Goto Github PK

 top_relevant_doc_pids = [embedding2pid[x] for y in I for x in y]

total_embeddings.shape= torch.Size([596847651, 128])
len(embedding2pid)= 596847471

I:
[[406822639 117817889 368295592 ... 340840928  44116417  58849792]
 [ 47088131 301290898  47088044 ... 126308015 284803824 288093606]
 [ 47088131 442551456  47088044 ... 540478083 300737160 316488359]
 ...
 [150980019 102136117 410337895 ...  39484138 289150555 571997339]
 [404093947 376114314 404093855 ... 166449253  32960475 549894389]
 [421370365  89232576 399436776 ... 501989308 415117993 378354340]]

IndexError occurred. Details:
Current y: [406822639 117817889 368295592 ... 340840928  44116417  58849792]
Current x: 596867001

Traceback (most recent call last):
  File "retrieve.py", line 140, in <module>
    raise e
  File "retrieve.py", line 138, in <module>
    raise e
  File "retrieve.py", line 132, in <module>
    pid = embedding2pid[x]
IndexError: list index out of range

wget https://msmarco.z22.web.core.windows.net/msmarcoranking/collectionandqueries.tar.gz
wget https://msmarco.z22.web.core.windows.net/msmarcoranking/triples.train.small.tar.gz
wget https://msmarco.z22.web.core.windows.net/msmarcoranking/top1000.eval.tar.gz

@staticmethod
def collate_fn(samples,tokenizer):

    def trim_padding(input_ids,padding_id):
        ## because we padding it to make length in the preprocess script
        ## we need to trim the padded sequences in a 2-dimensional tensor to the length of the longest non-padded sequence
        non_pad_mask = input_ids != padding_id
        non_pad_lengths = non_pad_mask.sum(dim=1)
        max_length = non_pad_lengths.max().item()
        trimmed_tensor = input_ids[:,:max_length]
        return trimmed_tensor

    queries  = [x[0] for x in samples]
    pos_docs = [x[1] for x in samples]
    neg_docs = [x[2] for x in samples]

    query_input_ids = torch.from_numpy(np.stack(queries).astype(np.int32))
    query_attention_mask = (query_input_ids != tokenizer.mask_token_id).int() ## not pad token, called *query augmentation* in the paper

    doc_input_ids = torch.from_numpy(np.stack(pos_docs+neg_docs).astype(np.int32))
    doc_input_ids = trim_padding(doc_input_ids,padding_id = tokenizer.pad_token_id)
    doc_attetion_mask = (doc_input_ids != tokenizer.pad_token_id).int()