TrainBatch

class lightning_ir.data.data.TrainBatch(queries: Sequence[str], docs: Sequence[Sequence[str]], query_ids: Sequence[str] | None = None, doc_ids: Sequence[Sequence[str]] | None = None, qrels: List[Dict[str, int]] | None = None, targets: Tensor | None = None)[source]

Bases: RankBatch

A batch of ranking data that combines multiple RankSample instances

Parameters:

  • queries (Sequence[str]) – List of query texts

  • docs (Sequence[Sequence[str]]) – List of list of document texts

  • query_ids (Sequence[str], optional) – Optional list of query ids

  • doc_ids (Sequence[Sequence[str]], optional) – Optional list of list of document ids

  • qrels (List[Dict[str, Any]], optional) – Optional list of dictionaries mapping document ids to relevance labels

  • targets (torch.Tensor, optional) – Optional list of target labels denoting the relevane of a document for the query

__init__(queries: Sequence[str], docs: Sequence[Sequence[str]], query_ids: Sequence[str] | None = None, doc_ids: Sequence[Sequence[str]] | None = None, qrels: List[Dict[str, int]] | None = None, targets: Tensor | None = None) → None

Methods

__init__(queries, docs[, query_ids, ...])

Attributes

doc_ids

qrels

query_ids

targets

queries

docs