PyTerrier Integration

API of ir_axioms.integrations.pyterrier

See the integrations guide and the walkthroughs for usage examples.

ir_axioms.integrations.pyterrier.transformers

KwikSortReranker dataclass

KwikSortReranker(
    *,
    axiom: Axiom[Query, Document],
    pivot_selection: PivotSelection = RandomPivotSelection(),
    text_field: str | None = "text",
    verbose: bool = False,
)

Bases: Transformer

axiom instance-attribute

axiom: Axiom[Query, Document]

pivot_selection class-attribute instance-attribute

pivot_selection: PivotSelection = RandomPivotSelection()

text_field class-attribute instance-attribute

text_field: str | None = 'text'

verbose class-attribute instance-attribute

verbose: bool = False

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/transformers.py

def transform(self, inp: DataFrame) -> DataFrame:
    require_columns(inp, {"qid", "docno"})
    inp = ensure_query_columns_hashable(inp)
    query_cols = list(query_columns(inp))
    query_rankings = inp.groupby(
        by=query_cols,
        group_keys=True,
        sort=False,
    )
    if len(query_rankings) == 0:
        return inp
    return concat(
        [
            self._transform_group(
                group_keys=dict(zip(query_cols, grouping)),
                res=ranking,
            )
            for grouping, ranking in tqdm(
                query_rankings,
                desc="KwikSort re-rank",
                unit="query",
                disable=not self.verbose,
            )
        ]
    )

AggregatedAxiomaticPreferences dataclass

AggregatedAxiomaticPreferences(
    axioms: Sequence[Axiom[Query, Document]],
    aggregations: Sequence[
        Callable[[Sequence[float]], float]
    ] = (lambda: [max, min, mean])(),
    text_field: str | None = "text",
    verbose: bool = False,
)

Bases: Transformer

axioms instance-attribute

axioms: Sequence[Axiom[Query, Document]]

aggregations class-attribute instance-attribute

aggregations: Sequence[
    Callable[[Sequence[float]], float]
] = field(default_factory=lambda: [max, min, mean])

text_field class-attribute instance-attribute

text_field: str | None = 'text'

verbose class-attribute instance-attribute

verbose: bool = False

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/transformers.py

def transform(self, inp: DataFrame) -> DataFrame:
    require_columns(inp, {"qid", "docno"})
    inp = ensure_query_columns_hashable(inp)
    query_cols = list(query_columns(inp))
    query_rankings = inp.groupby(
        by=query_cols,
        group_keys=True,
        sort=False,
    )
    if len(query_rankings) == 0:
        inp["features"] = None
        return inp
    return concat(
        [
            self._transform_group(
                group_keys=dict(zip(query_cols, grouping)),
                res=ranking,
            )
            for grouping, ranking in tqdm(
                query_rankings,
                desc="Aggregate axiom preferences",
                unit="query",
                disable=not self.verbose,
            )
        ]
    )

AxiomaticPreferences dataclass

AxiomaticPreferences(
    axioms: Sequence[Axiom],
    axiom_names: Sequence[str] | None = None,
    text_field: str | None = "text",
    verbose: bool = False,
)

Bases: Transformer

axioms instance-attribute

axioms: Sequence[Axiom]

axiom_names class-attribute instance-attribute

axiom_names: Sequence[str] | None = None

text_field class-attribute instance-attribute

text_field: str | None = 'text'

verbose class-attribute instance-attribute

verbose: bool = False

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/transformers.py

def transform(self, inp: DataFrame) -> DataFrame:
    require_columns(inp, {"qid", "docno"})
    inp = ensure_query_columns_hashable(inp)
    query_cols = list(query_columns(inp))
    query_rankings = inp.groupby(
        by=query_cols,
        group_keys=True,
        sort=False,
    )
    if len(query_rankings) == 0:
        inp = inp.merge(
            inp,
            on=query_cols,
            suffixes=("_a", "_b"),
            sort=False,
        )
        for axiom_name in self._axiom_names:
            inp[f"{axiom_name}_preference"] = None
        return inp
    return concat(
        [
            self._transform_group(
                group_keys=dict(zip(query_cols, grouping)),
                res=ranking,
            )
            for grouping, ranking in tqdm(
                query_rankings,
                desc="Compute axiom preferences",
                unit="query",
                disable=not self.verbose,
                maxinterval=0.2,
            )
        ]
    )

ir_axioms.integrations.pyterrier.estimator

EstimatorKwikSortReranker dataclass

EstimatorKwikSortReranker(
    *,
    axioms: Sequence[Axiom[Query, Document]],
    target: Axiom[Query, Document] = ORACLE(),
    estimator: ScikitLearnEstimator,
    pivot_selection: PivotSelection = RandomPivotSelection(),
    text_field: str | None = "text",
    verbose: bool = False,
)

Bases: Estimator

A KwikSort re-ranker that uses a scikit-learn estimator to estimate the target preferences (e.g., ORACLE preferences) from the preferences of multiple axioms. This class extends from PyTerrier's Estimator so that it can be fitted/trained like any other PyTerrier estimator.

name class-attribute instance-attribute

name = 'EstimatorKwikSortReranker'

axioms instance-attribute

axioms: Sequence[Axiom[Query, Document]]

target class-attribute instance-attribute

target: Axiom[Query, Document] = field(
    default_factory=ORACLE
)

estimator instance-attribute

estimator: ScikitLearnEstimator

pivot_selection class-attribute instance-attribute

pivot_selection: PivotSelection = RandomPivotSelection()

text_field class-attribute instance-attribute

text_field: str | None = 'text'

verbose class-attribute instance-attribute

verbose: bool = False

fit

fit(
    topics_or_res_tr: DataFrame,
    qrels_tr: DataFrame,
    topics_or_res_va: DataFrame | None = None,
    qrels_va: DataFrame | None = None,
) -> Transformer

Train the model with the given rankings and qrels.

Source code in ir_axioms/integrations/pyterrier/estimator.py

def fit(
    self,
    topics_or_res_tr: DataFrame,
    qrels_tr: DataFrame,
    topics_or_res_va: DataFrame | None = None,
    qrels_va: DataFrame | None = None,
) -> Transformer:
    """
    Train the model with the given rankings and qrels.
    """

    # Check if validation data is provided.
    if topics_or_res_va is not None or qrels_va is not None:
        warn("Validation data will be ignored in this estimator.", UserWarning)

    # Check if the required columns are present.
    require_columns(topics_or_res_tr, {"qid", "docno"})
    require_columns(qrels_tr, {"qid", "docno", "label"})

    # Check if there are any results to fit to.
    if len(topics_or_res_tr.index) == 0:
        raise ValueError("No results to fit to.")

    # Merge the rankings with the qrels to have the labels available.
    topics_or_res_tr = topics_or_res_tr.merge(
        qrels_tr,
        on=["qid", "docno"],
        how="left",
        suffixes=(
            "#old#",
            "",
        ),  # Keep the old columns to warn about overwriting.
    )
    # If the columns were overwritten, warn the user.
    overwritten_columns = [
        col.removesuffix("#old#")
        for col in topics_or_res_tr.columns
        if col.endswith("#old#")
    ]
    if len(overwritten_columns) > 0:
        warn(
            f"Columns were overwritten by the qrels: {', '.join(overwritten_columns)}.",
            UserWarning,
        )
        topics_or_res_tr = topics_or_res_tr.drop(
            columns=[f"{col}#old#" for col in overwritten_columns]
        )

    # Extract the queries and documents from the updated ranking.
    queries = load_queries(topics_or_res_tr)
    documents = load_documents(
        topics_or_res_tr,
        text_column=self.text_field,
    )

    # Group into batches of common queries and documents for each query.
    query_document_pairs = zip(queries, documents)
    query_document_batches = [
        (
            query,
            [document for _, document in query_documents],
        )
        for query, query_documents in groupby(
            query_document_pairs, key=lambda x: x[0]
        )
    ]

    # Using the grouped queries and documents, fit the estimator axiom to the given target axiom.
    self._estimator_axiom.fit(
        target=self.target,
        inputs_outputs=query_document_batches,
    )

    # Return the fitted re-ranker (i.e., this object).
    return self

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/estimator.py

def transform(self, inp: DataFrame) -> DataFrame:
    # Apply KwikSort re-ranking with the internal re-ranker.
    # The internal re-ranker uses an estimator axiom, which is assumed to be fitted at this point.
    return self._reranker.transform(inp)

ir_axioms.integrations.pyterrier.experiment

AxiomaticExperiment dataclass

AxiomaticExperiment(
    *,
    retrieval_systems: Sequence[Transformer],
    names: Sequence[str] | None = None,
    topics: DataFrame,
    qrels: DataFrame,
    axioms: Sequence[Axiom],
    axiom_names: Sequence[str] | None = None,
    depth: int | None = 10,
    filter_by_qrels: bool = True,
    filter_by_topics: bool = False,
    text_field: str | None = "text",
    parallel_jobs: int = 1,
    parallel_backend: Literal["joblib", "ray"] = "joblib",
    verbose: bool = False,
)

retrieval_systems instance-attribute

retrieval_systems: Sequence[Transformer]

names class-attribute instance-attribute

names: Sequence[str] | None = None

topics instance-attribute

topics: DataFrame

qrels instance-attribute

qrels: DataFrame

axioms instance-attribute

axioms: Sequence[Axiom]

axiom_names class-attribute instance-attribute

axiom_names: Sequence[str] | None = None

depth class-attribute instance-attribute

depth: int | None = 10

filter_by_qrels class-attribute instance-attribute

filter_by_qrels: bool = True

filter_by_topics class-attribute instance-attribute

filter_by_topics: bool = False

text_field class-attribute instance-attribute

text_field: str | None = 'text'

parallel_jobs class-attribute instance-attribute

parallel_jobs: int = 1

parallel_backend class-attribute instance-attribute

parallel_backend: Literal['joblib', 'ray'] = 'joblib'

verbose class-attribute instance-attribute

verbose: bool = False

preferences cached property

preferences: DataFrame

Return a dataframe with each axiom's preference for each document pair, query and retrieval system.

The returned dataframe has the following columns: - name: Retrieval system name (if found in source) - qid: Query ID - query: Query text - docno_a: Document A ID - score_a: Document A score - rank_a: Document A rank - label_a: Document A relevance label (if qrels were given) - docno_b: Document B ID - score_b: Document B score - rank_b: Document B rank - rank_b: Document B relevance label (if qrels were given) - ORIG_preference: Preference from original retrieved ranking - ORACLE_preference: Preference from qrels (NaN if qrels are missing) - _preference: Preference from each axiom

preference_distribution cached property

preference_distribution: DataFrame

Return a dataframe with each axiom's preference distribution across all documents, queries and retrieval systems.

The returned dataframe has the following columns: - axiom: Axiom name - axiom == 0: Absolute number of preferences equal to 0. - axiom == ORIG: Absolute number of preferences equal to the original ranking preference. - axiom != ORIG: Absolute number of preferences opposing the original ranking preference.

preference_consistency cached property

preference_consistency: DataFrame

Return a dataframe with each axiom's consistency with the original ranking preferences and oracle preferences across all documents, queries and retrieval systems.

The returned dataframe has the following columns: - axiom: Axiom name - ORIG_consistency: Relative consistency of non-zero preferences with ORIG preferences. - ORACLE_consistency: Relative consistency of non-zero preferences with ORACLE preferences. - _consistency: Relative consistency of non-zero preferences with ORIG preferences from system (if found in source)

inconsistent_pairs cached property

inconsistent_pairs: DataFrame

ir_axioms.integrations.pyterrier.utils

FilterTopicsTransformer dataclass

FilterTopicsTransformer(topics: DataFrame)

Bases: Transformer

Retain only queries that are contained in the topics.

topics instance-attribute

topics: DataFrame

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/utils.py

def transform(self, inp: DataFrame) -> DataFrame:
    return inp[inp["qid"].isin(self.topics["qid"])]

FilterQrelsTransformer dataclass

FilterQrelsTransformer(qrels: DataFrame)

Bases: Transformer

Retain only query-document pairs that are contained in the qrels.

qrels instance-attribute

qrels: DataFrame

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/utils.py

def transform(self, inp: DataFrame) -> DataFrame:
    return inp[
        inp["qid"].isin(self.qrels["qid"])
        & inp["docno"].isin(self.qrels["docno"])
    ]

JoinQrelsTransformer dataclass

JoinQrelsTransformer(qrels: DataFrame)

Bases: Transformer

Join query-document pairs with their relevance labels.

qrels instance-attribute

qrels: DataFrame

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/utils.py

def transform(self, inp: DataFrame) -> DataFrame:
    qrels = self.qrels
    require_columns(qrels, {"qid", "docno", "label"})
    return inp.merge(self.qrels, on=["qid", "docno"], how="left")

AddNameTransformer dataclass

AddNameTransformer(system_name: str)

Bases: Transformer

Add a colum containing the system's name.

system_name instance-attribute

system_name: str

transform

transform(inp: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/utils.py

def transform(self, inp: DataFrame) -> DataFrame:
    inp["name"] = self.system_name
    return inp

inject_pyterrier

inject_pyterrier(
    index_location: TerrierIndex
    | Path
    | str
    | Any
    | None = None,
    text_field: str | None = "text",
    tokeniser: TerrierTokeniser | Any | None = None,
    dataset: Dataset | str | None = None,
    injector: Injector = _default_injector,
) -> None

Source code in ir_axioms/integrations/pyterrier/utils.py

def inject_pyterrier(
    index_location: TerrierIndex | Path | str | Any | None = None,
    text_field: str | None = "text",
    tokeniser: TerrierTokeniser | Any | None = None,
    dataset: Dataset | str | None = None,
    injector: Injector = _default_injector,
) -> None:
    injector.binder.bind(
        interface=TermTokenizer,
        to=InstanceProvider(
            TerrierTermTokenizer(
                tokeniser=tokeniser,
                index_location=index_location,
            )
        ),
        scope=singleton,
    )

    if index_location is not None:
        injector.binder.bind(
            interface=IndexStatistics,
            # Note: Wrap in a provider, as otherwise the injector's type check accesses the (lazy) properties and starts the JVM.
            to=InstanceProvider(
                TerrierIndexStatistics(index_location=index_location)
            ),
            scope=singleton,
        )
        injector.binder.bind(
            interface=TextStatistics[Document],
            to=InstanceProvider(
                TerrierDocumentTextStatistics(index_location=index_location)
            ),
            scope=singleton,
        )

        if text_field is not None:
            injector.binder.bind(
                interface=TextContents[Document],
                to=InstanceProvider(
                    TerrierDocumentTextContents(
                        index_location=index_location,
                        text_field=text_field,
                    )
                ),
                scope=singleton,
            )

    if dataset is not None:
        injector.binder.bind(
            interface=TextContents[Query],
            to=InstanceProvider(
                IrdsQueryTextContents(
                    dataset=dataset,
                )
            ),
            scope=singleton,
        )
        injector.binder.bind(
            interface=TextContents[Document],
            to=InstanceProvider(
                IrdsDocumentTextContents(
                    dataset=dataset,
                )
            ),
            scope=singleton,
        )

    reset_binding_scopes(injector)

require_columns

require_columns(
    ranking: DataFrame, expected_columns: set[str]
) -> None

Source code in ir_axioms/integrations/pyterrier/utils.py

def require_columns(
    ranking: DataFrame,
    expected_columns: set[str],
) -> None:
    columns: set[str] = set(ranking.columns)
    missing_columns: set[str] = expected_columns - columns
    if len(missing_columns) > 0:
        raise ValueError(
            f"Expected columns "
            f"{', '.join(expected_columns)} but got columns "
            f"{', '.join(columns)} (missing columns "
            f"{', '.join(missing_columns)})."
        )

query_columns

query_columns(
    df: DataFrame, qid: bool = True
) -> Sequence[str]

Return the columns that are used to identify a query in the ranking.

Note: This function patches the query_columns function from PyTerrier to also include the query_toks column. Remove this patch, once https://github.com/terrier-org/pyterrier/issues/566 is resolved.

Source code in ir_axioms/integrations/pyterrier/utils.py

def query_columns(df: DataFrame, qid: bool = True) -> Sequence[str]:
    """
    Return the columns that are used to identify a query in the ranking.

    Note: This function patches the `query_columns` function from PyTerrier to also include the `query_toks` column.
    Remove this patch, once https://github.com/terrier-org/pyterrier/issues/566 is resolved.
    """
    columns = _query_columns(df, qid=qid)
    if "query_toks" in df.columns:
        columns = [*columns, "query_toks"]
    return columns

ensure_query_columns_hashable

ensure_query_columns_hashable(df: DataFrame) -> DataFrame

Source code in ir_axioms/integrations/pyterrier/utils.py

def ensure_query_columns_hashable(df: DataFrame) -> DataFrame:
    for column in df.columns:
        if all(isinstance(value, dict) for value in df[column]):
            df[column] = df[column].map(frozendict)
    return df

load_document

load_document(
    row: Series | Mapping[Hashable, Any],
    text_column: str | None = "text",
) -> Document

Source code in ir_axioms/integrations/pyterrier/utils.py

def load_document(
    row: Series | Mapping[Hashable, Any],
    text_column: str | None = "text",
) -> Document:
    text = (
        str(row[text_column])
        if text_column is not None and text_column in row.keys()
        else None
    )
    if "toks" in row.keys():
        tokens: Mapping[str, int] = row["toks"]
        if text is not None:
            text = TokenizedString(
                value=text,
                tokens=tokens,
            )
    return Document(
        id=str(row["docno"]),
        text=text,
        score=float(row["score"]) if "score" in row.keys() else None,
        rank=int(row["rank"]) if "rank" in row.keys() else None,
        relevance=float(row["label"]) if "label" in row.keys() else None,
    )

load_documents

load_documents(
    ranking: DataFrame, text_column: str | None = "text"
) -> Sequence[Document]

Source code in ir_axioms/integrations/pyterrier/utils.py

def load_documents(
    ranking: DataFrame,
    text_column: str | None = "text",
) -> Sequence[Document]:
    return [
        load_document(
            row=row,
            text_column=text_column,
        )
        for _, row in ranking.iterrows()
    ]

load_query

load_query(row: Series | Mapping[Hashable, Any]) -> Query

Source code in ir_axioms/integrations/pyterrier/utils.py

def load_query(row: Series | Mapping[Hashable, Any]) -> Query:
    text = str(row["query"]) if "query" in row.keys() else None
    if "query_toks" in row.keys():
        tokens: Mapping[str, int] = row["query_toks"]
        if text is not None:
            text = TokenizedString(
                value=text,
                tokens=tokens,
            )
    return Query(
        id=str(row["qid"]),
        text=text,
    )

load_queries

load_queries(ranking: DataFrame) -> Sequence[Query]

Source code in ir_axioms/integrations/pyterrier/utils.py

def load_queries(ranking: DataFrame) -> Sequence[Query]:
    return [load_query(row=row) for _, row in ranking.iterrows()]