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)