# AUTOGENERATED! DO NOT EDIT! File to edit: ../01_module_passage.ipynb.

# %% auto 0
__all__ = ['sample_dict_items', 'save_data', 'load_data', 'PassageData', 'sample_data', 'create_basic_search_package']

# %% ../01_module_passage.ipynb 4
import ir_datasets 
import json
import requests
import os 
import pandas as pd

from fastcore.utils import patch, patch_to
from typing import Dict, Tuple, Optional
from random import sample
from vespa.package import (
    ApplicationPackage, 
    Field, 
    FieldSet, 
    RankProfile, 
    QueryField
)

# %% ../01_module_passage.ipynb 7
def sample_dict_items(
    d: Dict, # dict to be samples from.                     
    n: int   # Number of samples 
) -> Dict: # dict with sampled values
    "Sample items from a dict."
    n = min(len(d), n)
    return {
        k: d[k]
        for k in sample(
            k=n, population=sorted(d)
        )
    }

# %% ../01_module_passage.ipynb 16
def save_data(
    corpus: Dict,  # Document corpus, see usage example below. 
    train_qrels: Dict,  # Training relevance scores, see usage example below.  
    train_queries: Dict,  # Training queries, see usage example below.   
    dev_qrels: Dict,  #  Development relevance scores, see usage example below.  
    dev_queries: Dict,  # Development queries, see usage example below.   
    file_path: str = "passage_sample.json"  # Valid JSON file path.
) -> None:  # Side-effect: data is saved to `file_path`.
    """Save data to disk.
    
    The main goal is to save sample data to disk.
    """
    data = {
        "corpus": corpus, 
        "train_qrels": train_qrels, 
        "train_queries": train_queries, 
        "dev_qrels": dev_qrels, 
        "dev_queries": dev_queries
    }
    with open(file_path, "w") as f:
        json.dump(data, f)
    
    

# %% ../01_module_passage.ipynb 20
def load_data(
    file_path: Optional[str] = None  # valid JSON file path contain data saved by `save_data`. If `None`, a pre-generated sample will be downloaded.
) -> Dict:  # See usage example below for expected format.  
    """Load data.
    
    The main goal is to load sample data from disk.
    If a `file_path` is not provided, a pre-generated data sample 
    will be downloaded.
    """
    if not file_path:
        file_url = "https://data.vespa.oath.cloud/learntorank/passage/passage_sample.json"
        data_json = json.loads(
            requests.get(file_url).text
        )
    else:
        with open(file_path, "r") as f:
            data_json = json.load(f)
        
    return {
        "corpus": data_json["corpus"], 
        "train_qrels": data_json["train_qrels"], 
        "train_queries": data_json["train_queries"],
        "dev_qrels": data_json["dev_qrels"],
        "dev_queries": data_json["dev_queries"]
    }
    

# %% ../01_module_passage.ipynb 30
class PassageData(object):
    def __init__(
        self, 
        corpus: Optional[Dict] = None,  # Document corpus, see usage example below.
        train_qrels: Optional[Dict] = None,  # Training relevance scores, see usage example below.  
        train_queries: Optional[Dict] = None,  # Training queries, see usage example below.   
        dev_qrels: Optional[Dict] = None,  #  Development relevance scores, see usage example below.  
        dev_queries: Optional[Dict] = None,  # Development queries, see usage example below.   
    ):
        "Container for passage data"
        self.corpus = corpus
        self.train_qrels = train_qrels
        self.train_queries = train_queries
        self.dev_qrels = dev_qrels
        self.dev_queries = dev_queries
    
    def __repr__(self):
        return "PassageData({}, {}, {}, {}, {})".format(
            "corpus" if self.corpus else None,
            "train_qrels" if self.train_qrels else None,
            "train_queries" if self.train_queries else None,
            "dev_qrels" if self.dev_qrels else None,
            "dev_queries" if self.dev_queries else None,            
        )
    
    def __eq__(self, obj: "PassageData"):
        return (self.corpus == obj.corpus and 
                self.train_qrels == obj.train_qrels and 
                self.train_queries == obj.train_queries and 
                self.dev_qrels == obj.dev_qrels and 
                self.dev_queries == obj.dev_queries)

# %% ../01_module_passage.ipynb 35
@patch
def save(self: PassageData, file_path: str = 'passage_sample.json'):
    save_data(
        corpus=self.corpus, 
        train_qrels=self.train_qrels, 
        train_queries=self.train_queries, 
        dev_qrels=self.dev_qrels, 
        dev_queries=self.dev_queries,
        file_path=file_path
    )

# %% ../01_module_passage.ipynb 39
@patch_to(PassageData, cls_method=True)
def load(
    cls, file_path: Optional[str] = None  # valid JSON file path contain data saved by save_data. If None, a pre-generated sample will be downloaded.
) -> PassageData:
    "Load passage data from disk."
    data = load_data(file_path=file_path)
    return PassageData(
        corpus=data.get("corpus", None), 
        train_qrels=data.get("train_qrels", None), 
        train_queries=data.get("train_queries", None),
        dev_qrels=data.get("dev_qrels", None),
        dev_queries=data.get("dev_queries", None),        
    )
    

# %% ../01_module_passage.ipynb 43
@patch(as_prop=True)
def summary(self: PassageData):
    "Summary of the size of the dataset components."
    print(
f"Number of documents: {len(self.corpus)}\n\
Number of train queries: {len(self.train_queries)}\n\
Number of train relevance judgments: {len(self.train_qrels)}\n\
Number of dev queries: {len(self.dev_queries)}\n\
Number of dev relevance judgments: {len(self.dev_qrels)}"
    )

# %% ../01_module_passage.ipynb 46
@patch
def get_corpus(self: PassageData) -> pd.DataFrame:
    return pd.DataFrame([{"doc_id": k, "text": v} for k, v in self.corpus.items()])

# %% ../01_module_passage.ipynb 49
@patch
def get_queries(
    self: PassageData, 
    type: str  # Either 'train' or 'dev'.
) -> pd.DataFrame:  # DataFrame conaining 'query_id' and 'query'.
    "Get query data."
    assert type in ['train', 'dev'], "'type' should be either 'train' or 'dev'."
    if type == 'train':
        data = self.train_queries
    elif type == 'dev':
        data = self.dev_queries
    return pd.DataFrame([{"query_id": k, "query": v} for k, v in data.items()])

# %% ../01_module_passage.ipynb 56
@patch
def get_labels(
    self: PassageData, 
    type: str  # Either 'train' or 'dev'.
) -> Dict:  # pyvespa-formatted labeled data 
    "Get labeled data"
    assert type in ['train', 'dev'], "'type' should be either 'train' or 'dev'."
    if type == 'train':
        queries = self.train_queries
        qrels = self.train_qrels
    elif type == 'dev':
        queries = self.dev_queries
        qrels = self.dev_qrels        
    return [
        {
            "query_id": query_id, 
            "query": query, 
            "relevant_docs": [{"id": doc_id, "score": score} for doc_id, score in qrels[query_id].items()]
        } 
        for query_id, query in queries.items()
    ]

# %% ../01_module_passage.ipynb 62
def sample_data(
    n_relevant: int,  # The number of relevant documents to sample.
    n_irrelevant: int,  # The number of non-judged documents to sample. 
) -> PassageData:  

    """
    Sample data from the passage ranking dataset.
    
    The final sample contains `n_relevant` train relevant documents, 
    `n_relevant` dev relevant documents and `n_irrelevant` random documents
    sampled from the entire corpus. 
    
    All the relevant sampled documents, both from train and dev sets, 
    are guaranteed to be on the `corpus_sample`, which will contain 
    2 * `n_relevant` + `n_irrelevant` documents.
    """
    passage_corpus = ir_datasets.load("msmarco-passage")
    passage_train = ir_datasets.load("msmarco-passage/train/judged")    
    passage_dev = ir_datasets.load("msmarco-passage/dev/judged")
    
    train_qrels_dict = passage_train.qrels_dict()
    train_qrels_sample = sample_dict_items(d=train_qrels_dict, n=n_relevant)

    dev_qrels_dict = passage_dev.qrels_dict() 
    dev_qrels_sample = sample_dict_items(d=dev_qrels_dict, n=n_relevant)
    
    train_queries = {k:v for k,v in passage_train.queries_iter()}
    train_queries_sample = {k: train_queries[k] for k in train_qrels_sample.keys()}
    dev_queries = {k:v for k,v in passage_dev.queries_iter()}
    dev_queries_sample = {k: dev_queries[k] for k in dev_qrels_sample.keys()}

    train_positive_doc_id_samples = [
        doc_id[0]
        for doc_id in [list(docs.keys()) for docs in train_qrels_sample.values()]
    ]
    dev_positive_doc_id_samples = [
        doc_id[0]
        for doc_id in [list(docs.keys()) for docs in dev_qrels_sample.values()]
    ]

    docs_id = [x[0] for x in passage_corpus.docs_iter()]    
    docs_store = passage_corpus.docs_store()
    
    negative_doc_id_samples = sample(
        k=n_irrelevant, population=sorted(docs_id)
    )
    doc_id_samples = list(
        set(
            train_positive_doc_id_samples
            + dev_positive_doc_id_samples
            + negative_doc_id_samples
        )
    )
    corpus_sample = {k: docs_store.get(k)[1] for k in doc_id_samples}

    return PassageData(
        corpus=corpus_sample,
        train_qrels=train_qrels_sample,
        train_queries=train_queries_sample,
        dev_qrels=dev_qrels_sample,
        dev_queries=dev_queries_sample,
    )


# %% ../01_module_passage.ipynb 80
def create_basic_search_package(
    name: str="PassageRanking"  # Name of the application
) -> ApplicationPackage: # pyvespa [ApplicationPackage](https://pyvespa.readthedocs.io/en/latest/reference-api.html#applicationpackage) instance.
    """
    Create a basic Vespa application package for passage ranking.
    
    *Vespa fields*:
        
    The application contain two string fields: `doc_id` and `text`.
    
    *Vespa rank functions*:
    
    The application contain two rank profiles: 
    [bm25](https://docs.vespa.ai/en/reference/bm25.html) and
    [nativeRank](https://docs.vespa.ai/en/reference/nativerank.html).
    
    """
    app_package = ApplicationPackage(name=name)
    app_package.schema.add_fields(
        Field(name="doc_id", type="string", indexing=["attribute", "summary"]),
        Field(name="text", type="string", indexing=["index", "summary"], index="enable-bm25"),
    )
    app_package.schema.add_field_set(
        FieldSet(name="default", fields=["text"])
    )
    app_package.schema.add_rank_profile(
        RankProfile(
            name="bm25",
            first_phase="bm25(text)",
            summary_features=["bm25(text)"]
        )
    )
    app_package.schema.add_rank_profile(
        RankProfile(
            name="native_rank", 
            first_phase="nativeRank(text)"
        )
    )
    app_package.query_profile.add_fields(
        QueryField(name="maxHits", value=10000)
    )
    return app_package
