<jupyter_start><jupyter_text>Pyserini integrationUsing pyserini to retrieve the top 1000 documents using BM25 and rank using ColBERT.<jupyter_code>INDEX_PATH = "/home/kenzosaki/mestrado/data/pt_msmarco/indexes"
QUERIES_PATH = "/home/kenzosaki/mestrado/data/pt_msmarco/google_translate/portuguese.queries.dev.small.tsv"
QRELS_PATH = "/home/kenzosaki/mestrado/data/en_msmarco/data/msmarco_ans_small/qrels.dev.small.tsv"
CHECKPOINT_PATH = "/home/kenzosaki/mestrado/repos/ColBERT/experiments/ColBertimbau/train.py/colbertimbau/checkpoints/colbert-200000.dnn"
BM25_PARAMETERS = {
    "k1": 1.25, 
    "b": 0.74
}

RM3_PARAMETERS = {
    "fb_docs": 10, 
    "fb_terms": 30, 
    "original_query_weight": 0.4
}
COLBERT_PARAMETERS = {
    "model_name": "neuralmind/bert-base-portuguese-cased",
    "query_maxlen": 32,
    "doc_maxlen": 180,
    "dim": 128,
    "similarity_metric": "cosine",
    "mask_punctuation": True
}
TOPK = 1000<jupyter_output><empty_output><jupyter_text>Loading indexes of MsMARCO-PT<jupyter_code>from pyserini.index import IndexReader
from pyserini.analysis import  get_lucene_analyzer
index_reader = IndexReader(INDEX_PATH)
index_reader.stats()<jupyter_output><empty_output><jupyter_text>Preparing the BM25 retrieval<jupyter_code>from pyserini.search import SimpleSearcher
from typing import List
searcher = SimpleSearcher(INDEX_PATH)
analyser = get_lucene_analyzer("pt")
searcher.set_analyzer(analyser) # para textos em portugues
searcher.set_bm25(**BM25_PARAMETERS)
searcher.set_rm3(**RM3_PARAMETERS)
query = "o que pode detectar o exame de urina"
hits = searcher.search(query, TOPK)
index_reader.analyze(query, analyser)
# TODO: da pra adicionar mais stopwords aqui. do NLTK por ex?
analyser.stopwordSet.toString()<jupyter_output><empty_output><jupyter_text>Loading queries<jupyter_code>import pandas as pd

queries_df = pd.read_csv(QUERIES_PATH, sep='\t', header=None, names=['query_id', 'query'])
qrels_df = pd.read_csv(QRELS_PATH, sep='\t', header=None, names=['query_id', 'x', 'pid', 'y'])
queries_df = queries_df.set_index('query_id').loc[qrels_df["query_id"].unique()]
queries_df.info()
queries_df.head()<jupyter_output><empty_output><jupyter_text>BM25 retrieval<jupyter_code>from tqdm import tqdm

def run_bm25(queries_df: pd.DataFrame, searcher: SimpleSearcher, output_path: str, k: int = 1000) -> None:

    # Output file
    with open(output_path, "w") as f:
        for query_id, row in tqdm(queries_df.iterrows(), total=len(queries_df) , desc="- Running BM25"):
            # BM25 retrieval
            query = row["query"]
            hits = searcher.search(query, k)
            
            for rank, hit in enumerate(hits, start=1):
                f.write("{}\t{}\t{}\n".format(query_id, hit.docid, rank))
run_bm25(queries_df, searcher, "data/runs/bm25_rm3_run.tsv", k=TOPK)<jupyter_output>- Running BM25: 100%|██████████| 105/105 [00:18<00:00,  5.76it/s]<jupyter_text>Loading ColBERT<jupyter_code>from colbert.modeling.inference import ModelInference
from colbert.modeling.colbert import ColBERT
from colbert.utils.utils import load_checkpoint
import torch
def load_colbert_from_checkpoint(path_to_checkpoint: str, device: torch.device):
    """
    Carrega o modelo ColBERT a partir checkpoint.
    """
    # Parametros usados para treino
    colbert = ColBERT.from_pretrained("neuralmind/bert-base-portuguese-cased", **COLBERT_PARAMETERS)
    colbert.to(device)

    checkpoint = load_checkpoint(path_to_checkpoint, colbert, do_print=True)

    colbert.eval()

    return colbert, checkpoint
device = torch.device("cuda")
# Preparing to inference
colbert, checkpoint = load_colbert_from_checkpoint(CHECKPOINT_PATH, device)
# colbert.to(device)
inference_model = ModelInference(colbert)<jupyter_output>Some weights of the model checkpoint at neuralmind/bert-base-portuguese-cased were not used when initializing ColBERT: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.seq_relationship.weight', 'cls.predictions.transform.LayerNorm.weight']
- This IS expected if you are initializing ColBERT from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).
- This IS NOT expected if you are initializing ColBERT from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).
Some weights of ColBERT were not initialized from the model checkpoint at neuralmind/bert-base-portuguese-cased and are newly initialized: ['[...]<jupyter_text>Evaluating ColBERT<jupyter_code>import numpy as np
import torch

def get_docs_and_pids_from_hits(hits):

    docs = [hit.contents for hit in hits]
    pids = [hit.docid for hit in hits]

    return docs, pids

def get_scores(inference_model, query, docs, bs):
    
    with torch.no_grad():
        Q = inference_model.queryFromText([query]).squeeze(0)
        D_ = inference_model.docFromText(docs, bsize=bs)
        scores = inference_model.colbert.score(Q, D_).cpu()

    return scores


def run_colbert(queries_df: pd.DataFrame,
                searcher: SimpleSearcher,
                inference_model: ModelInference, 
                output_path: str, 
                bs: int = 32,
                k: int = 1000) -> None:

    # Output file
    with open(output_path, "w") as f:
        for query_id, row in tqdm(queries_df.iterrows(), total=len(queries_df) , desc="- Running ColBERT"):
            # BM25 retrieval
            query = row["query"]
            hits = searcher.search(query, k)

            docs, pids = get_docs_and_pids_from_hits(hits)
            scores = get_scores(inference_model, query, docs, bs)

            sorted_indexes = torch.argsort(scores, descending=True)
            
            for rank, index in enumerate(sorted_indexes, start=1):
                f.write("{}\t{}\t{}\n".format(query_id, pids[index], rank))
checkpoint_name = CHECKPOINT_PATH.split("/")[-1]
checkpoint_name
run_colbert(queries_df, searcher, inference_model, f"data/runs/{checkpoint_name}.run.tsv", bs=512, k=TOPK)<jupyter_output>- Running ColBERT: 100%|██████████| 105/105 [06:48<00:00,  3.89s/it]
{'fraction_non_alphanumeric': 0.11524572146333804, 'fraction_numerical': 0.019155283403988067, 'mean_word_length': 4.529513888888889, 'pattern_counts': {'":': 11, '<': 28, '<?xml version=': 0, '>': 28, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '599245', 'license': 'MIT', 'n_tokens_mistral': 2339, 'n_tokens_neox': 2195, 'n_words': 456, 'path': 'pyserini_integration.ipynb', 'repo_name': 'kenzosakiyama/ColBERT', 'size': '6369'}