# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# This notebook experiments with the MSMarco Dataset, nothing that exciting to see here right now. But fun to toy with.
#
# # Setup & Indexing
#
# The next several cells we setup the SolrClient (run everytime), we download the MSMarco dataset, and index it into Solr.

# +
from ltr.client import SolrClient
from ltr.index import rebuild


client=SolrClient()
# -

# Download MSMarco Corpus and Queries -> data directory
from ltr import download
msmarco_resources = ['https://msmarco.blob.core.windows.net/msmarcoranking/msmarco-docs.tsv.gz',
                     'https://msmarco.blob.core.windows.net/msmarcoranking/msmarco-docs-lookup.tsv.gz',
                     'https://msmarco.blob.core.windows.net/msmarcoranking/msmarco-doctrain-qrels.tsv.gz',
                     'https://msmarco.blob.core.windows.net/msmarcoranking/msmarco-doctrain-queries.tsv.gz']
download(msmarco_resources, dest='data/')

# +
# Index to Solr. Docs & Question fro convenience

import csv
import sys
import gzip

csv.field_size_limit(sys.maxsize)

def marco_docs():
    with gzip.open('data/msmarco-docs.tsv.gz', 'rt') as tsvfile:
        reader = csv.reader(tsvfile, delimiter='\t')
        i = 0;
        for row in reader:
        
                yield {"id": row[0],
                       "url": row[1],
                       "title": row[2],
                       "body": row[3],
                       "type": "document"}
                i+=1
                if i % 10000 == 0:
                    print("Dumped (%s/%s) %s" % (i, 3213835, row[1]))
                    
def marco_questions():
    with gzip.open('data/msmarco-doctrain-queries.tsv.gz', 'rt') as tsvfile:
        reader = csv.reader(tsvfile, delimiter='\t')
        i = 0;
        for row in reader:
            yield {"id": 'Q' + row[0],
                   "url": '',
                   "title": row[1],
                   "body": row[1],
                   "type": "question"}
            i+=1
            if i % 10000 == 0:
                print("Dumped (q) (%s/%s) %s" % (i, 367013, row[1]))
                    

def marco_questions_and_docs():
    for q in marco_questions():
        yield q
    for d in marco_docs():
        yield d
        


# +
from ltr.client import SolrClient
from ltr.index import rebuild


client=SolrClient()
rebuild(client, index='msmarco', doc_type='passage', doc_src=marco_questions_and_docs())
# -

# # Some experiments

# +
# Parse out all the msmarco qrels

from ltr.helpers.msmarco.evaluate import QRel

qrels = {}

for qrel in QRel.read_qrels():
    qrels[qrel.qid] = qrel

# +
import random

sample_qids = [k for k in qrels.keys()]
random.shuffle(sample_qids)
sample_qids = sample_qids[:100]
sample_qids

# +
from ltr.helpers.solr_escape import esc_kw


def eval_one(client, qrel, params, at=50):
    """ Eval a single Solr query param for MRR"""
    kw = esc_kw(qrel.keywords)
    
    import copy
    params = copy.deepcopy(params)

    params['q'] = kw
    params['rows'] = at
    params['fl'] = 'id'
    if 'fq' in params:
        params['fq'] = [params['fq'], 'type:document']
    else:
        params['fq'] = "type:document"
    
    hits = client.query(index='msmarco', query=params)
    
    ranking = [hit['id'] for hit in hits]
    rr = qrel.eval_rr(ranking)
    return rr

def eval_many(client, qrels, sample_qids, params):
    """ Execute many Solr searches, return MRR
        (this should show distribution...)"""
    sum_rr = 0.0
    num_evald = 0
    
    HIST_SLOTS = 100
    
    all_rrs = []

    print("Running %s queries" % len(sample_qids))

    for qid in sample_qids:
        qrel = qrels[qid]        
        rr = eval_one(client, qrel, params)
        
        all_rrs.append(rr)
  
        sum_rr += rr
        num_evald += 1
        print("%s, last rr: %s, mrr: %s" % (num_evald, rr, sum_rr / num_evald))      
    return all_rrs


# +
all_rrs_title = eval_many(client, qrels, sample_qids, {"qf": "title", "defType": "edismax"})
all_rrs_body = eval_many(client, qrels, sample_qids, {"qf": "body", "defType": "edismax"})
all_rrs_all = eval_many(client, qrels, sample_qids, {"qf": "title body url", "defType": "edismax"})
all_rrs_all_tie = eval_many(client, qrels, sample_qids, {"qf": "title body url", "tie": 1.0, "defType": "edismax"})
all_rrs_all_pf2 = eval_many(client, qrels, sample_qids, {"qf": "body", "pf2": "body", "tie": 1.0,  "defType": "edismax"})



# +
# Matplotlib and seaborn for plotting
import matplotlib.pyplot as plt
import matplotlib
import numpy as np

# %matplotlib inline

matplotlib.rcParams['font.size'] = 18
matplotlib.rcParams['figure.dpi'] = 200

from IPython.core.pylabtools import figsize
figsize(15, 5)


def draw_rr_dist(all_rrs, experiment_name):
    bins = [0.0, 0.05,0.1,0.15,0.2,0.25,
                 0.3,0.35,0.4,0.45,0.5,
                 0.55,0.6,0.65,0.7,0.75,
                 0.8,0.85,0.9,0.95,1.0]
    
    all_rrs = np.array(all_rrs)


    # Make the histogram using matplotlib, bins must be integet
    plt.hist(all_rrs, color = 'blue', edgecolor = 'black',
             bins = bins)

    # Add labels
    plt.title('RR by Query )' + experiment_name)
    plt.xlabel('RR'); plt.ylabel('Num Qs');
    plt.axvline(all_rrs.mean(), color='k', linestyle='dashed', linewidth=1)
    
# -



draw_rr_dist(all_rrs_title, "Title Search")

draw_rr_dist(all_rrs_body, "Body Search")

draw_rr_dist(all_rrs_all, "Title Body Url Dismax Search")

draw_rr_dist(all_rrs_all_tie, "Title Body Url Dismax Search, tie")

# +
draw_rr_dist(all_rrs_all_pf2, "PF2 Search")


# -





StarcoderdataJupyterScriptsDedupFiltered