File size: 920 Bytes
cd6775d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 | from colbert.infra import Run, RunConfig, ColBERTConfig
from colbert import Searcher
import torch
searcher = Searcher(index=f"/data1/liuyaoyang/Papers/icml2025/multi_rag/RAG/Search-in-the-Chain/ColBERT/experiments/webq/indexes/webq.nbits=2")
query = "where is rome italy located on a map?"
# mask = torch.tensor([[1, 1, 0, 0, 1, 1, 0, 0, 0, 1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
# 0, 0, 0, 0, 0, 0, 0, 0]], dtype = torch.int32)
mask = torch.tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
0, 0, 0, 0, 0, 0, 0, 0]], dtype = torch.int32)
pids1, ranks, scores = searcher.search_with_mask(query, mask, k=100)
# ["where", "is" ,"rome" ,"italy", "located", "on", "a" ,"map","?"]
pids2, ranks, scores = searcher.search_with_sentence(query, ["where is rome italy", "where is located on", "a map","?"], k=100)
print(pids1[:3])
print(pids2[:3])
print(pids1 == pids2) |