<jupyter_start><jupyter_text>LTR - Learning-To-Rank Reranking Models for MS MARCO Passage  Install Dependencies<jupyter_code>from google.colab import auth
auth.authenticate_user()
import os
#%%capture
!curl -O https://download.java.net/java/GA/jdk11/9/GPL/openjdk-11.0.2_linux-x64_bin.tar.gz
!mv openjdk-11.0.2_linux-x64_bin.tar.gz /usr/lib/jvm/; cd /usr/lib/jvm/; tar -zxvf openjdk-11.0.2_linux-x64_bin.tar.gz
!update-alternatives --install /usr/bin/java java /usr/lib/jvm/jdk-11.0.2/bin/java 1
!update-alternatives --set java /usr/lib/jvm/jdk-11.0.2/bin/java
os.environ["JAVA_HOME"] = "/usr/lib/jvm/jdk-11.0.2"
!gsutil cp gs://luanps/information_retrieval/pyserini/bm25/run.msmarco-passage.bm25tuned.txt .
!git clone https://github.com/luanps/pyserini.git  --recurse-submodules
!pip install pyserini 
!pip install -U pip setuptools wheel
!pip install -U spacy
!python -m spacy download en_core_web_sm
!pip install faiss-cpu
'''!wget https://rgw.cs.uwaterloo.ca/JIMMYLIN-bucket0/pyserini-models/model-ltr-ibm.tar.gz -P collections/msmarco-ltr-passage/
!tar -xzvf collections/msmarco-ltr-passage/model-ltr-ibm.tar.gz -C collections/msmarco-ltr-passage/

!wget https://rgw.cs.uwaterloo.ca/JIMMYLIN-bucket0/pyserini-models/model-ltr-msmarco-passage-mrr-v1.tar.gz -P runs/
!tar -xzvf runs/model-ltr-msmarco-passage-mrr-v1.tar.gz -C runs'''
!wget https://msmarco.blob.core.windows.net/msmarcoranking/qidpidtriples.train.full.2.tsv.gz -P collections/msmarco-passage/	
!gzip -d collections/msmarco-passage/qidpidtriples.train.full.2.tsv.gz
#Download MSMARCO dataset

!mkdir collections/msmarco-passage
!wget https://msmarco.blob.core.windows.net/msmarcoranking/collectionandqueries.tar.gz -P collections/msmarco-passage
# Alternative mirror:
# wget https://www.dropbox.com/s/9f54jg2f71ray3b/collectionandqueries.tar.gz -P collections/msmarco-passage
!tar xvfz collections/msmarco-passage/collectionandqueries.tar.gz -C collections/msmarco-passage
!cp -r pyserini/tools/ .
!cp -r pyserini/scripts .
'''!mkdir collections/msmarco-ltr-passage/

!python pyserini/scripts/ltr_msmarco/convert_queries.py \
  --input collections/msmarco-passage/queries.eval.small.tsv \
  --output collections/msmarco-ltr-passage/queries.eval.small.json 

!python pyserini/scripts/ltr_msmarco/convert_queries.py \
  --input collections/msmarco-passage/queries.dev.small.tsv \
  --output collections/msmarco-ltr-passage/queries.dev.small.json

!python pyserini/scripts/ltr_msmarco/convert_queries.py \
  --input collections/msmarco-passage/queries.train.tsv \
  --output collections/msmarco-ltr-passage/queries.train.json

#!gsutil cp -r collections/msmarco-ltr-passage/  gs://luanps/information_retrieval/pyserini/ltr/'''
!gsutil cp -r   gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ collections/<jupyter_output>Copying gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ibm_model/body/output.t1.5.bin...
Copying gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ibm_model/body/source.vcb...
Copying gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ibm_model/body/target.vcb...
Copying gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ibm_model/text_bert_tok/output.t1.5.bin...
-
==> NOTE: You are performing a sequence of gsutil operations that may
run significantly faster if you instead use gsutil -m cp ... Please
see the -m section under "gsutil help options" for further information
about when gsutil -m can be advantageous.

Copying gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ibm_model/text_bert_tok/source.vcb...
Copying gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ibm_model/text_bert_tok/target.vcb...
Copying gs://luanps/information_retrieval/pyserini/ltr/msmarco-ltr-passage/ibm_model/[...]<jupyter_text>Run LTR inference<jupyter_code>!python -m pyserini.search.lucene.ltr  --input run.msmarco-passage.bm25tuned.txt \
  --input-format tsv \
  --model runs/msmarco-passage-ltr-mrr-v1 \
  --index msmarco-passage-ltr \
  --data passage \
  --ibm-model collections/msmarco-ltr-passage/ibm_model/ \
  --queries collections/msmarco-ltr-passage \
--output run.ltr.msmarco-passage.tsv 
#MRR Eval
!python -m pyserini.eval.msmarco_passage_eval msmarco-passage-dev-subset run.ltr.msmarco-passage.tsv >> ltr_mrr_eval.txt


#TREC Eval
!python -m pyserini.eval.convert_msmarco_run_to_trec_run --input run.ltr.msmarco-passage.tsv --output run.ltr.msmarco-passage.trec

!python -m pyserini.eval.trec_eval -c -mrecall.1000 \
                                      -mmap msmarco-passage-dev-subset \
                                      run.ltr.msmarco-passage.trec >> ltr_trec_eval.txt
!python -m pyserini.eval.convert_msmarco_run_to_trec_run \
   

!python -m pyserini/tools/scripts/msmarco/convert_msmarco_to_trec_qrels \
   --input tools/topics-and-qrels/qrels.msmarco-passage.dev-subset.txt --output collections/msmarco-passage/qrels.dev.small.trec
$ tools/eval/trec_eval.9.0.4/trec_eval -c -mrecall.1000 -mmap \
   collections/msmarco-passage/qrels.dev.small.trec runs/run.ltr.msmarco-passage.trec
map                   	all	0.2551
recall_1000           	all	0.8573<jupyter_output><empty_output><jupyter_text>LTR Feature Extraction<jupyter_code>!python -m pyserini.search.lucene.ltr  --input run.msmarco-passage.bm25tuned.txt \
  --input-format tsv \
  --model runs/msmarco-passage-ltr-mrr-v1 \
  --index msmarco-passage-ltr \
  --data passage \
  --ibm-model collections/msmarco-ltr-passage/ibm_model/ \
  --queries collections/msmarco-ltr-passage \
--output run.ltr.msmarco-passage.tsv 
import sys

# We're going to explicitly use a local installation of Pyserini (as opposed to a pip-installed one).
# Comment these lines out to use a pip-installed one instead.
sys.path.insert(0, './')

import argparse
import json
import multiprocessing
import os
import pickle
import time

import numpy as np
import pandas as pd
from tqdm import tqdm
from collections import defaultdict
#from pyserini.search.lucene.ltr._search_msmarco import MsmarcoLtrSearcher
#from pyserini.search.lucene.ltr import *


def train_data_loader(task='triple', neg_sample=20, random_seed=12345):
    print(f'train_{task}_sampled_with_{neg_sample}_{random_seed}.pickle')
    if os.path.exists(f'./collections/msmarco-ltr-passage/train_{task}_sampled_with_{neg_sample}_{random_seed}.pickle'):
        sampled_train = pd.read_pickle(f'./collections/msmarco-ltr-passage/train_{task}_sampled_with_{neg_sample}_{random_seed}.pickle')
        print(sampled_train.shape)
        print(sampled_train.index.get_level_values('qid').drop_duplicates().shape)
        print(sampled_train.groupby('qid').count().mean())
        print(sampled_train.head(10))
        print(sampled_train.info())
        return sampled_train
    else:
        if task == 'triple':
            train = pd.read_csv('./collections/msmarco-passage/qidpidtriples.train.full.2.tsv', sep="\t",
                                names=['qid', 'pos_pid', 'neg_pid'], dtype=str)#np.int32)
            pos_half = train[['qid', 'pos_pid']].rename(columns={"pos_pid": "pid"}).drop_duplicates()
            pos_half['rel'] = np.int32(1)
            neg_half = train[['qid', 'neg_pid']].rename(columns={"neg_pid": "pid"}).drop_duplicates()
            neg_half['rel'] = np.int32(0)
            del train
            sampled_neg_half = []
            for qid, group in tqdm(neg_half.groupby('qid')):
                sampled_neg_half.append(group.sample(n=min(neg_sample, len(group)), random_state=random_seed))
            sampled_train = pd.concat([pos_half] + sampled_neg_half, axis=0, ignore_index=True)
            sampled_train = sampled_train.sort_values(['qid', 'pid']).set_index(['qid', 'pid'])
            print(sampled_train.shape)
            print(sampled_train.index.get_level_values('qid').drop_duplicates().shape)
            print(sampled_train.groupby('qid').count().mean())
            print(sampled_train.head(10))
            print(sampled_train.info())

            sampled_train.to_pickle(f'./collections/msmarco-ltr-passage/train_{task}_sampled_with_{neg_sample}_{random_seed}.pickle')
        elif task == 'rank':
            qrel = defaultdict(list)
            with open("./collections/msmarco-passage/qrels.train.tsv") as f:
                for line in f:
                    topicid, _, docid, rel = line.strip().split('\t')
                    assert rel == "1", line.split(' ')
                    qrel[topicid].append(docid)

            qid2pos = defaultdict(list)
            qid2neg = defaultdict(list)
            with open("./runs/msmarco-passage/run.train.small.tsv") as f:
                for line in tqdm(f):
                    topicid, docid, rank = line.split()
                    assert topicid in qrel
                    if docid in qrel[topicid]:
                        qid2pos[topicid].append(docid)
                    else:
                        qid2neg[topicid].append(docid)
            sampled_train = []
            for topicid, pos_list in tqdm(qid2pos.items()):
                neg_list = random.sample(qid2neg[topicid], min(len(qid2neg[topicid]), neg_sample))
                for positive_docid in pos_list:
                    sampled_train.append((int(topicid), int(positive_docid), 1))
                for negative_docid in neg_list:
                    sampled_train.append((int(topicid), int(negative_docid), 0))
            sampled_train = pd.DataFrame(sampled_train, columns=['qid', 'pid', 'rel'], dtype=np.int32)
            sampled_train = sampled_train.sort_values(['qid', 'pid']).set_index(['qid', 'pid'])
            print(sampled_train.shape)
            print(sampled_train.index.get_level_values('qid').drop_duplicates().shape)
            print(sampled_train.groupby('qid').count().mean())
            print(sampled_train.head(10))
            print(sampled_train.info())

            sampled_train.to_pickle(f'./collections/msmarco-ltr-passage/train_{task}_sampled_with_{neg_sample}_{random_seed}.pickle')
        else:
            raise Exception('unknown parameters')
        return sampled_train


"""
Running prediction on candidates
"""
def dev_data_loader(file, format, data, top=100):
    if format == 'tsv':
        dev = pd.read_csv(file, sep="\t",
                          names=['qid', 'pid', 'rank'],
                          dtype={'qid': 'S','pid': 'S', 'rank':'i',})
    elif format == 'trec':
        dev = pd.read_csv(file, sep="\s+",
                    names=['qid', 'q0', 'pid', 'rank', 'score', 'tag'],
                    usecols=['qid', 'pid', 'rank'],
                    dtype={'qid': 'S','pid': 'S', 'rank':'i',})
    else:
        raise Exception('unknown parameters')
    assert dev['qid'].dtype == object
    assert dev['pid'].dtype == object
    assert dev['rank'].dtype == np.int32
    dev = dev[dev['rank']<=top]
    if data == 'passage':
        dev_qrel = pd.read_csv('tools/topics-and-qrels/qrels.msmarco-passage.dev-subset.txt', sep=" ",
                            names=["qid", "q0", "pid", "rel"], usecols=['qid', 'pid', 'rel'],
                            dtype={'qid': 'S','pid': 'S', 'rel':'i'})
    elif data == 'document':
        dev_qrel = pd.read_csv('tools/topics-and-qrels/qrels.msmarco-doc.dev.txt', sep="\t",
                            names=["qid", "q0", "pid", "rel"], usecols=['qid', 'pid', 'rel'],
                            dtype={'qid': 'S','pid': 'S', 'rel':'i'})
    assert dev['qid'].dtype == object
    assert dev['pid'].dtype == object
    assert dev['rank'].dtype == np.int32
    dev = dev.merge(dev_qrel, left_on=['qid', 'pid'], right_on=['qid', 'pid'], how='left')
    dev['rel'] = dev['rel'].fillna(0).astype(np.int32)
    dev = dev.sort_values(['qid', 'pid']).set_index(['qid', 'pid'])

    print(dev.shape)
    print(dev.index.get_level_values('qid').drop_duplicates().shape)
    print(dev.groupby('qid').count().mean())
    print(dev.head(10))
    print(dev.info())

    dev_rel_num = dev_qrel[dev_qrel['rel'] > 0].groupby('qid').count()['rel']

    recall_point = [10, 20, 50, 100, 200, 500, 1000, 2000, 5000, 10000]
    recall_curve = {k: [] for k in recall_point}
    for qid, group in tqdm(dev.groupby('qid')):
        group = group.reset_index()
        assert len(group['pid'].tolist()) == len(set(group['pid'].tolist()))
        total_rel = dev_rel_num.loc[qid]
        query_recall = [0 for k in recall_point]
        for t in group.sort_values('rank').itertuples():
            if t.rel > 0:
                for i, p in enumerate(recall_point):
                    if t.rank <= p:
                        query_recall[i] += 1
        for i, p in enumerate(recall_point):
            if total_rel > 0:
                recall_curve[p].append(query_recall[i] / total_rel)
            else:
                recall_curve[p].append(0.)

    for k, v in recall_curve.items():
        avg = np.mean(v)
        print(f'recall@{k}:{avg}')

    return dev, dev_qrel


def query_loader():
    queries = {}
    with open(f"{args['queries']}") as f:
        for line in f:
            query = json.loads(line)
            qid = query.pop('id')
            query['analyzed'] = query['analyzed'].split(" ")
            query['text'] = query['text_unlemm'].split(" ")
            query['text_unlemm'] = query['text_unlemm'].split(" ")
            query['text_bert_tok'] = query['text_bert_tok'].split(" ")
            queries[qid] = query
    return queries


def eval_mrr(dev_data):
    score_tie_counter = 0
    score_tie_query = set()
    MRR = []
    for qid, group in tqdm(dev_data.groupby('qid')):
        group = group.reset_index()
        rank = 0
        prev_score = None
        assert len(group['pid'].tolist()) == len(set(group['pid'].tolist()))
        # stable sort is also used in LightGBM

        for t in group.sort_values('score', ascending=False, kind='mergesort').itertuples():
            if prev_score is not None and abs(t.score - prev_score) < 1e-8:
                score_tie_counter += 1
                score_tie_query.add(qid)
            prev_score = t.score
            rank += 1
            if t.rel > 0:
                MRR.append(1.0 / rank)
                break
            elif rank == 10 or rank == len(group):
                MRR.append(0.)
                break

    score_tie = f'score_tie occurs {score_tie_counter} times in {len(score_tie_query)} queries'
    print(score_tie)
    mrr_10 = np.mean(MRR).item()
    print(f'MRR@10:{mrr_10} with {len(MRR)} queries')
    return {'score_tie': score_tie, 'mrr_10': mrr_10}


def eval_recall(dev_qrel, dev_data):
    dev_rel_num = dev_qrel[dev_qrel['rel'] > 0].groupby('qid').count()['rel']

    score_tie_counter = 0
    score_tie_query = set()

    recall_point = [10,20,50,100,200,250,300,333,400,500,1000]
    recall_curve = {k: [] for k in recall_point}
    for qid, group in tqdm(dev_data.groupby('qid')):
        group = group.reset_index()
        rank = 0
        prev_score = None
        assert len(group['pid'].tolist()) == len(set(group['pid'].tolist()))
        # stable sort is also used in LightGBM
        total_rel = dev_rel_num.loc[qid]
        query_recall = [0 for k in recall_point]
        for t in group.sort_values('score', ascending=False, kind='mergesort').itertuples():
            if prev_score is not None and abs(t.score - prev_score) < 1e-8:
                score_tie_counter += 1
                score_tie_query.add(qid)
            prev_score = t.score
            rank += 1
            if t.rel > 0:
                for i, p in enumerate(recall_point):
                    if rank <= p:
                        query_recall[i] += 1
        for i, p in enumerate(recall_point):
            if total_rel > 0:
                recall_curve[p].append(query_recall[i] / total_rel)
            else:
                recall_curve[p].append(0.)

    score_tie = f'score_tie occurs {score_tie_counter} times in {len(score_tie_query)} queries'
    print(score_tie)
    res = {'score_tie': score_tie}

    for k, v in recall_curve.items():
        avg = np.mean(v)
        print(f'recall@{k}:{avg}')
        res[f'recall@{k}'] = avg

    return res
args = {
    'input': 'run.msmarco-passage.bm25tuned.txt',
    'input_format': 'tsv',
    'model': 'runs/msmarco-passage-ltr-mrr-v1',
    'index': 'msmarco-passage-ltr',
    'data': 'passage',
    'ibm_model': 'collections/msmarco-ltr-passage/ibm_model/',
    'queries':  'collections/msmarco-ltr-passage/queries.train.json',
    'output': 'run.ltr.msmarco-passage.tsv',
    'reranking_top': 10000,
    'max_passage': True,
    'neg_sample': 10
  }
import pickle

#sampled_train = train_data_loader(task='triple', neg_sample = args['neg_sample'])
#!gsutil cp collections/msmarco-ltr-passage/train_triple_sampled_with_10_12345.pickle gs://luanps/information_retrieval/pyserini/ltr/
!gsutil cp   gs://luanps/information_retrieval/pyserini/ltr/train_triple_sampled_with_10_12345.pickle  collections/msmarco-ltr-passage/
sampled_train = pickle.load(open('collections/msmarco-ltr-passage/train_triple_sampled_with_10_12345.pickle','rb'))

sampled_train = sampled_train.reset_index()
sampled_train['qid'] = sampled_train['qid'].astype(str)
sampled_train['pid'] = sampled_train['pid'].astype(str)
queries = query_loader()
#dev, dev_qrel = dev_data_loader(args['input'], args['input_format'], args['data'], args['reranking_top'])
from pyserini.search.lucene.ltr import MsmarcoLtrSearcher

searcher = MsmarcoLtrSearcher(args['model'], args['ibm_model'], args['index'], args['data'])
searcher.add_fe()
def batch_extract(searcher, df, queries, fe):
    tasks = []
    task_infos = []
    group_lst = []

    for qid, group in tqdm(df.groupby('qid')):
        '''try:
            query_dict = queries[str(qid)]

        except:
            continue'''
            
        task = {
            "qid": qid,
            "docIds": [],
            "rels": [],
            "query_dict": queries[str(qid)]
        }
        for t in group.reset_index().itertuples():
            if (searcher.data == 'document'):
                if (searcher.index_reader.doc(t.pid) != None):
                    task["docIds"].append(t.pid)
                    task_infos.append((qid, t.pid, t.rel))
            else:
                task["docIds"].append(t.pid)
                task_infos.append((qid, t.pid, t.rel))
        tasks.append(task)
        group_lst.append((qid, len(task['docIds'])))


        if len(tasks) == 1000:
            features = fe.batch_extract(tasks)
            task_infos = pd.DataFrame(task_infos, columns=['qid', 'pid', 'rel'])
            group = pd.DataFrame(group_lst, columns=['qid', 'count'])
            print(features.shape)
            print(task_infos.qid.drop_duplicates().shape)
            print(group.mean())
            print(features.head(10))
            print(features.info())
            yield task_infos, features, group
            tasks = []
            task_infos = []
            group_lst = []
    # deal with rest
    if len(tasks) > 0:
        features = fe.batch_extract(tasks)
        task_infos = pd.DataFrame(task_infos, columns=['qid', 'pid', 'rel'])
        group = pd.DataFrame(group_lst, columns=['qid', 'count'])
        print(features.shape)
        print(task_infos.qid.drop_duplicates().shape)
        print(group.mean())
        print(features.head(10))
        print(features.info())
        yield task_infos, features, group

    return
batch_info = []
for extracted in batch_extract(searcher, sampled_train, queries, searcher.fe): #searcher.batch_extract(sampled_train, queries, searcher.fe): #
    batch_info.append(extracted)
pickle.dump(batch_info,open('msmarco_passage_train.ltr.features.pkl','wb'))
!gsutil cp msmarco_passage_train.ltr.features.pkl gs://luanps/information_retrieval/pyserini/ltr/<jupyter_output><empty_output><jupyter_text>Save train data<jupyter_code>!gsutil cp gs://luanps/information_retrieval/pyserini/ltr/train_triple_sampled_with_10_12345.pickle .
!head -n2 collections/msmarco-passage/queries.train.tsv
import pickle 
train_triple = pickle.load(open('train_triple_sampled_with_10_12345.pickle','rb'))
train_triple = train_triple.reset_index()
import pandas as pd

collection = pd.read_csv('./collections/msmarco-passage/collection.tsv', sep="\t",
                          names=['pid', 'passage'],
                          dtype={'pid': 'i', 'passage':'S',})

queries = pd.read_csv('./collections/msmarco-passage/queries.train.tsv', sep="\t",
                          names=['qid', 'query'],
                          dtype={'qid': 'i', 'query':'S',})
train_query = train_triple.merge(queries, on='qid')
train_collection = train_query.merge(collection, on='pid')
pickle.dump(train_collection,open('msmarco_passage_train.ltr.collection.pkl','wb'))
!gsutil cp msmarco_passage_train.ltr.collection.pkl gs://luanps/information_retrieval/pyserini/ltr/<jupyter_output>Copying file://msmarco_passage_train.ltr.collection.pkl [Content-Type=application/octet-stream]...
/ [0 files][    0.0 B/  1.2 GiB]                                                
==> NOTE: You are uploading one or more large file(s), which would run
significantly faster if you enable parallel composite uploads. This
feature can be enabled by editing the
"parallel_composite_upload_threshold" value in your .boto
configuration file. However, note that if you do this large files will
be uploaded as `composite objects
<https://cloud.google.com/storage/docs/composite-objects>`_,which
means that any user who downloads such objects will need to have a
compiled crcmod installed (see "gsutil help crcmod"). This is because
without a compiled crcmod, computing checksums on composite objects is
so slow that gsutil disables downloads of composite objects.

\
Operation completed over 1 objects/1.2 GiB.
{'fraction_non_alphanumeric': 0.14386803289987593, 'fraction_numerical': 0.013922712861278316, 'mean_word_length': 2.899659559218778, 'pattern_counts': {'":': 6, '<': 21, '<?xml version=': 0, '>': 30, 'https://': 8, 'lorem ipsum': 0, 'www.': 1, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 9, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '446613', 'license': 'Apache-2.0', 'n_tokens_mistral': 7933, 'n_tokens_neox': 7655, 'n_words': 1269, 'path': 'Pyserini_LTR.ipynb', 'repo_name': 'luanps/pyserini', 'size': '21763'}