Skip to content

Tools

API of ir_axioms.tools

See tools and dependency injection for the default tool bindings and how to replace them.

ir_axioms.tools.tokenizer

TermTokenizer

Bases: Protocol

terms

terms(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/base.py
def terms(self, text: str) -> Sequence[str]: ...

terms_unordered

terms_unordered(text: str) -> Collection[str]
Source code in ir_axioms/tools/tokenizer/base.py
def terms_unordered(self, text: str) -> Collection[str]:
    return self.terms(text)

unique_terms

unique_terms(text: str) -> Set[str]
Source code in ir_axioms/tools/tokenizer/base.py
def unique_terms(self, text: str) -> Set[str]:
    return set(self.terms(text))

iter_terms

iter_terms(texts: Iterable[str]) -> Iterator[Sequence[str]]
Source code in ir_axioms/tools/tokenizer/base.py
def iter_terms(self, texts: Iterable[str]) -> Iterator[Sequence[str]]:
    for text in texts:
        yield self.terms(text)

iter_terms_unordered

iter_terms_unordered(
    texts: Iterable[str],
) -> Iterator[Collection[str]]
Source code in ir_axioms/tools/tokenizer/base.py
def iter_terms_unordered(self, texts: Iterable[str]) -> Iterator[Collection[str]]:
    for text in texts:
        yield self.terms_unordered(text)

iter_unique_terms

iter_unique_terms(
    texts: Iterable[str],
) -> Iterator[Set[str]]
Source code in ir_axioms/tools/tokenizer/base.py
def iter_unique_terms(self, texts: Iterable[str]) -> Iterator[Set[str]]:
    for text in texts:
        yield self.unique_terms(text)

SentenceTokenizer

Bases: Protocol

sentences

sentences(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/base.py
def sentences(self, text: str) -> Sequence[str]: ...

iter_sentences

iter_sentences(
    texts: Iterable[str],
) -> Iterator[Sequence[str]]
Source code in ir_axioms/tools/tokenizer/base.py
def iter_sentences(self, texts: Iterable[str]) -> Iterator[Sequence[str]]:
    for text in texts:
        yield self.sentences(text)

NltkTermTokenizer dataclass

NltkTermTokenizer(
    *, lemmatize: bool = False, lowercase: bool = False
)

Bases: TermTokenizer

lemmatize class-attribute instance-attribute

lemmatize: bool = False

lowercase class-attribute instance-attribute

lowercase: bool = False

terms

terms(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/nltk.py
def terms(self, text: str) -> Sequence[str]:
    terms = word_tokenize(text)
    # Lowercase before lemmatizing, as the WordNet lemmatizer is case-sensitive.
    if self.lowercase:
        terms = [term.lower() for term in terms]
    if self.lemmatize:
        terms = [self._lemmatizer.lemmatize(term) for term in terms]
    return terms

NltkSentenceTokenizer

NltkSentenceTokenizer()

Bases: SentenceTokenizer

Source code in ir_axioms/tools/tokenizer/nltk.py
def __init__(self) -> None:
    download_nltk_dependencies("punkt")
    download_nltk_dependencies("punkt_tab")

sentences

sentences(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/nltk.py
def sentences(self, text: str) -> Sequence[str]:
    return sent_tokenize(text)

AnseriniTermTokenizer dataclass

AnseriniTermTokenizer(
    *, analyzer: Analyzer = default_analyzer()
)

Bases: TermTokenizer

analyzer class-attribute instance-attribute

analyzer: Analyzer = field(default_factory=default_analyzer)

terms

terms(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/pyserini.py
def terms(self, text: str) -> Sequence[str]:
    return self.analyzer.analyze(text)

TerrierTermTokenizer dataclass

TerrierTermTokenizer(
    *,
    tokeniser: TerrierTokeniser | Any | None = None,
    index_location: TerrierIndex
    | Path
    | str
    | Any
    | None = None,
)

Bases: TermTokenizer

tokeniser class-attribute instance-attribute

tokeniser: TerrierTokeniser | Any | None = None

index_location class-attribute instance-attribute

index_location: TerrierIndex | Path | str | Any | None = (
    None
)

terms

terms(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/pyterrier.py
@pt_java_required
def terms(self, text: str) -> Sequence[str]:
    # ADDITION: Is there a way to restore the *sequence* of tokens from a TokenizedString (i.e., that contains a mapping of tokens to frequencies but no position information)?

    # Tokenize the text with the configured tokenizer.
    reader = _J.StringReader(str(text))
    terms = [
        str(term)
        for term in self._tokeniser.tokenise(reader)  # type: ignore
        if term is not None
    ]

    # Post-process the text with the configured term pipeline.
    for pipeline in self._term_pipelines:
        terms = [
            str(term)
            for term in map(pipeline.pipelineTerm, terms)
            if term is not None
        ]
    return terms

terms_unordered

terms_unordered(text: str) -> Collection[str]
Source code in ir_axioms/tools/tokenizer/pyterrier.py
def terms_unordered(self, text: str) -> Collection[str]:
    # If the text is a TokenizedString (e.g., from TerrierTextContents), we can directly use its tokens.
    if isinstance(text, TokenizedString):
        return list(
            chain.from_iterable(
                repeat(term, count) for term, count in text.tokens.items()
            )
        )

    return super().terms_unordered(text)

unique_terms

unique_terms(text: str) -> Set[str]
Source code in ir_axioms/tools/tokenizer/pyterrier.py
def unique_terms(self, text: str) -> Set[str]:
    # If the text is a TokenizedString (e.g., from TerrierTextContents), we can directly use its tokens.
    if isinstance(text, TokenizedString):
        return set(text.tokens.keys())

    return super().unique_terms(text)

SpacyTermTokenizer dataclass

SpacyTermTokenizer(
    *,
    language_name: str = "en_core_web_sm",
    lemmatize: bool = True,
    lowercase: bool = True,
    remove_punctuation: bool = True,
    remove_stopwords: bool = True,
)

Bases: TermTokenizer

language_name class-attribute instance-attribute

language_name: str = 'en_core_web_sm'

lemmatize class-attribute instance-attribute

lemmatize: bool = True

lowercase class-attribute instance-attribute

lowercase: bool = True

remove_punctuation class-attribute instance-attribute

remove_punctuation: bool = True

remove_stopwords class-attribute instance-attribute

remove_stopwords: bool = True

terms

terms(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/spacy.py
def terms(self, text: str) -> Sequence[str]:
    return self._terms(self._language(text))

iter_terms

iter_terms(texts: Iterable[str]) -> Iterator[Sequence[str]]
Source code in ir_axioms/tools/tokenizer/spacy.py
def iter_terms(self, texts: Iterable[str]) -> Iterator[Sequence[str]]:
    for document in self._language.pipe(texts):
        yield self._terms(document)

iter_terms_unordered

iter_terms_unordered(
    texts: Iterable[str],
) -> Iterator[Collection[str]]
Source code in ir_axioms/tools/tokenizer/spacy.py
def iter_terms_unordered(self, texts: Iterable[str]) -> Iterator[Collection[str]]:
    return self.iter_terms(texts)

iter_unique_terms

iter_unique_terms(
    texts: Iterable[str],
) -> Iterator[Set[str]]
Source code in ir_axioms/tools/tokenizer/spacy.py
def iter_unique_terms(self, texts: Iterable[str]) -> Iterator[Set[str]]:
    for terms in self.iter_terms(texts):
        yield set(terms)

SpacySentenceTokenizer dataclass

SpacySentenceTokenizer(
    *, language_name: str = "en_core_web_sm"
)

Bases: SentenceTokenizer

language_name class-attribute instance-attribute

language_name: str = 'en_core_web_sm'

sentences

sentences(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/spacy.py
def sentences(self, text: str) -> Sequence[str]:
    document = self._language(text)
    return [sentence.text for sentence in document.sents]

iter_sentences

iter_sentences(
    texts: Iterable[str],
) -> Iterator[Sequence[str]]
Source code in ir_axioms/tools/tokenizer/spacy.py
def iter_sentences(self, texts: Iterable[str]) -> Iterator[Sequence[str]]:
    for document in self._language.pipe(texts):
        yield [sentence.text for sentence in document.sents]

BlingfireTermTokenizer dataclass

BlingfireTermTokenizer(*, remove_punctuation: bool = True)

Bases: TermTokenizer

remove_punctuation class-attribute instance-attribute

remove_punctuation: bool = True

terms

terms(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/blingfire.py
def terms(self, text: str) -> Sequence[str]:
    return [
        word
        for word in text_to_words(text).split()
        if not self.remove_punctuation or any(char.isalnum() for char in word)
    ]

BlingfireSentenceTokenizer dataclass

BlingfireSentenceTokenizer(
    *, language_name: str = "en_core_web_sm"
)

Bases: SentenceTokenizer

language_name class-attribute instance-attribute

language_name: str = 'en_core_web_sm'

sentences

sentences(text: str) -> Sequence[str]
Source code in ir_axioms/tools/tokenizer/blingfire.py
def sentences(self, text: str) -> Sequence[str]:
    return text_to_sentences(text).splitlines(keepends=False)

TokenizerModule

Bases: Module

configure

configure(binder: Binder) -> None
Source code in ir_axioms/tools/tokenizer/__init__.py
def configure(self, binder: Binder) -> None:
    binder.bind(
        interface=TermTokenizer,
        to=SpacyTermTokenizer,
        scope=singleton,
    )
    binder.bind(
        interface=SentenceTokenizer,
        to=SpacySentenceTokenizer,
        scope=singleton,
    )
    if is_blingfire_installed():
        binder.bind(
            interface=TermTokenizer,
            to=BlingfireTermTokenizer,
            scope=singleton,
        )
        binder.bind(
            interface=SentenceTokenizer,
            to=BlingfireSentenceTokenizer,
            scope=singleton,
        )

ir_axioms.tools.contents

TextContents

Bases: Generic[T], Protocol

contents

contents(input: T) -> str
Source code in ir_axioms/tools/contents/base.py
def contents(self, input: T) -> str: ...

iter_contents

iter_contents(inputs: Iterable[T]) -> Iterator[str]
Source code in ir_axioms/tools/contents/base.py
def iter_contents(self, inputs: Iterable[T]) -> Iterator[str]:
    for input in inputs:
        yield self.contents(input)

DocumentQueryTextContents dataclass

DocumentQueryTextContents(
    *,
    document_text_contents: TextContents[Document],
    query_text_contents: TextContents[Query],
)

Bases: TextContents[Document | Query]

document_text_contents instance-attribute

document_text_contents: TextContents[Document]

query_text_contents instance-attribute

query_text_contents: TextContents[Query]

contents

contents(input: Document | Query) -> str
Source code in ir_axioms/tools/contents/combine.py
def contents(self, input: Document | Query) -> str:
    if isinstance(input, Document):
        return self.document_text_contents.contents(input)
    elif isinstance(input, Query):
        return self.query_text_contents.contents(input)

GenerationInputOutputTextContents dataclass

GenerationInputOutputTextContents(
    *,
    generation_input_text_contents: TextContents[
        GenerationRequest
    ],
    generation_output_text_contents: TextContents[
        GenerationResponse
    ],
)

Bases: TextContents[GenerationRequest | GenerationResponse]

generation_input_text_contents instance-attribute

generation_input_text_contents: TextContents[
    GenerationRequest
]

generation_output_text_contents instance-attribute

generation_output_text_contents: TextContents[
    GenerationResponse
]

contents

contents(
    input: GenerationRequest | GenerationResponse,
) -> str
Source code in ir_axioms/tools/contents/combine.py
def contents(self, input: GenerationRequest | GenerationResponse) -> str:
    if isinstance(input, GenerationRequest):
        return self.generation_input_text_contents.contents(input)
    elif isinstance(input, GenerationResponse):
        return self.generation_output_text_contents.contents(input)

IrdsDocumentTextContents dataclass

IrdsDocumentTextContents(*, dataset: Dataset | str)

Bases: TextContents[Document]

dataset instance-attribute

dataset: Dataset | str

contents

contents(input: Document) -> str
Source code in ir_axioms/tools/contents/ir_datasets.py
def contents(self, input: Document) -> str:
    if input.text is not None:
        return input.text
    irds_document: GenericDoc = self._documents_store.get(input.id)
    return irds_document.default_text()

iter_contents

iter_contents(inputs: Iterable[Document]) -> Iterator[str]
Source code in ir_axioms/tools/contents/ir_datasets.py
def iter_contents(self, inputs: Iterable[Document]) -> Iterator[str]:
    inputs = list(inputs)
    missing_ids = {input.id for input in inputs if input.text is None}
    irds_documents: Mapping[str, GenericDoc] = (
        self._documents_store.get_many(missing_ids) if len(missing_ids) > 0 else {}
    )
    for input in inputs:
        if input.text is not None:
            yield input.text
        elif input.id in irds_documents:
            yield irds_documents[input.id].default_text()
        else:
            # Raise the same error as for a single missing document.
            yield self.contents(input)

IrdsQueryTextContents dataclass

IrdsQueryTextContents(*, dataset: Dataset | str)

Bases: TextContents[Query]

dataset instance-attribute

dataset: Dataset | str

contents

contents(input: Query) -> str
Source code in ir_axioms/tools/contents/ir_datasets.py
def contents(self, input: Query) -> str:
    if input.text is not None:
        return input.text
    return self._queries[input.id]

AnseriniDocumentTextContents dataclass

AnseriniDocumentTextContents(*, index_dir: Path | str)

Bases: TextContents[Document]

index_dir instance-attribute

index_dir: Path | str

contents

contents(input: Document) -> str
Source code in ir_axioms/tools/contents/pyserini.py
def contents(self, input: Document) -> str:
    if input.text is not None:
        return input.text
    document = self._searcher.doc(input.id)
    if document is None:
        raise KeyError(f"Document '{input.id}' not found in index.")
    return document.contents()

TerrierDocumentTextContents dataclass

TerrierDocumentTextContents(
    *,
    index_location: TerrierIndex | Path | str | Any,
    text_field: str = "text",
)

Bases: TextContents[Document]

index_location instance-attribute

index_location: TerrierIndex | Path | str | Any

text_field class-attribute instance-attribute

text_field: str = 'text'

contents

contents(input: Document) -> str
Source code in ir_axioms/tools/contents/pyterrier.py
def contents(self, input: Document) -> str:
    if input.text is not None:
        return input.text
    return self._document_contents(input.id)

HasText

Bases: Protocol

text property

text: str | None

SimpleTextContents

Bases: TextContents[HasText]

contents

contents(input: HasText) -> str
Source code in ir_axioms/tools/contents/simple.py
def contents(self, input: HasText) -> str:
    if input.text is not None:
        return input.text
    raise ValueError(f"Could not get text contents from: {input}")

ContentsModule

Bases: Module

configure

configure(binder: Binder) -> None
Source code in ir_axioms/tools/contents/__init__.py
def configure(self, binder: Binder) -> None:
    binder.bind(
        interface=TextContents[HasText],
        to=SimpleTextContents,  # type: ignore
        scope=singleton,
    )
    binder.bind(
        interface=TextContents[Query],
        to=SimpleTextContents,  # type: ignore
        scope=singleton,
    )
    binder.bind(
        interface=TextContents[Document],
        to=SimpleTextContents,  # type: ignore
        scope=singleton,
    )
    binder.bind(
        interface=TextContents[GenerationRequest],
        to=SimpleTextContents,  # type: ignore
        scope=singleton,
    )
    binder.bind(
        interface=TextContents[GenerationResponse],
        to=SimpleTextContents,  # type: ignore
        scope=singleton,
    )
    binder.bind(
        interface=TextContents[Document | Query],
        to=DocumentQueryTextContents,
        scope=singleton,
    )
    binder.bind(
        interface=TextContents[GenerationRequest | GenerationResponse],
        to=GenerationInputOutputTextContents,
        scope=singleton,
    )

ir_axioms.tools.text_statistics

TextStatistics

Bases: Protocol, Generic[T]

term_counts

term_counts(document: T) -> Mapping[str, int]
Source code in ir_axioms/tools/text_statistics/base.py
def term_counts(self, document: T) -> Mapping[str, int]: ...

term_count

term_count(document: T, term: str) -> int
Source code in ir_axioms/tools/text_statistics/base.py
def term_count(self, document: T, term: str) -> int:
    term_counts = self.term_counts(document)
    return term_counts.get(term, 0)

term_frequencies

term_frequencies(document: T) -> Mapping[str, float]
Source code in ir_axioms/tools/text_statistics/base.py
def term_frequencies(self, document: T) -> Mapping[str, float]:
    term_counts = self.term_counts(document)
    document_length = sum(term_counts.values())
    return {term: count / document_length for term, count in term_counts.items()}

term_frequency

term_frequency(document: T, term: str) -> float
Source code in ir_axioms/tools/text_statistics/base.py
def term_frequency(self, document: T, term: str) -> float:
    return self.term_frequencies(document).get(term, 0)

DocumentQueryTextStatistics dataclass

DocumentQueryTextStatistics(
    *,
    document_text_contents: TextStatistics[Document],
    query_text_contents: TextStatistics[Query],
)

Bases: TextStatistics[Document | Query]

document_text_contents instance-attribute

document_text_contents: TextStatistics[Document]

query_text_contents instance-attribute

query_text_contents: TextStatistics[Query]

term_counts

term_counts(
    document: Document | Query,
) -> Mapping[str, int]
Source code in ir_axioms/tools/text_statistics/combine.py
def term_counts(self, document: Document | Query) -> Mapping[str, int]:
    if isinstance(document, Document):
        return self.document_text_contents.term_counts(document)
    elif isinstance(document, Query):
        return self.query_text_contents.term_counts(document)

term_frequencies

term_frequencies(
    document: Document | Query,
) -> Mapping[str, float]
Source code in ir_axioms/tools/text_statistics/combine.py
def term_frequencies(self, document: Document | Query) -> Mapping[str, float]:
    if isinstance(document, Document):
        return self.document_text_contents.term_frequencies(document)
    elif isinstance(document, Query):
        return self.query_text_contents.term_frequencies(document)

GenerationInputOutputTextStatistics dataclass

GenerationInputOutputTextStatistics(
    *,
    generation_input_text_statistics: TextStatistics[
        GenerationRequest
    ],
    generation_output_text_statistics: TextStatistics[
        GenerationResponse
    ],
)

Bases: TextStatistics[GenerationRequest | GenerationResponse]

generation_input_text_statistics instance-attribute

generation_input_text_statistics: TextStatistics[
    GenerationRequest
]

generation_output_text_statistics instance-attribute

generation_output_text_statistics: TextStatistics[
    GenerationResponse
]

term_counts

term_counts(
    document: GenerationRequest | GenerationResponse,
) -> Mapping[str, int]
Source code in ir_axioms/tools/text_statistics/combine.py
def term_counts(
    self, document: GenerationRequest | GenerationResponse
) -> Mapping[str, int]:
    if isinstance(document, GenerationRequest):
        return self.generation_input_text_statistics.term_counts(document)
    elif isinstance(document, GenerationResponse):
        return self.generation_output_text_statistics.term_counts(document)

term_frequencies

term_frequencies(
    document: GenerationRequest | GenerationResponse,
) -> Mapping[str, float]
Source code in ir_axioms/tools/text_statistics/combine.py
def term_frequencies(
    self, document: GenerationRequest | GenerationResponse
) -> Mapping[str, float]:
    if isinstance(document, GenerationRequest):
        return self.generation_input_text_statistics.term_frequencies(document)
    elif isinstance(document, GenerationResponse):
        return self.generation_output_text_statistics.term_frequencies(document)

AnseriniTextStatistics dataclass

AnseriniTextStatistics(
    *,
    index_dir: Path | str,
    analyzer: Analyzer = default_analyzer(),
)

Bases: TextStatistics[DocumentType]

index_dir instance-attribute

index_dir: Path | str

analyzer class-attribute instance-attribute

analyzer: Analyzer = field(default_factory=default_analyzer)

term_counts

term_counts(document: DocumentType) -> Mapping[str, int]
Source code in ir_axioms/tools/text_statistics/pyserini.py
def term_counts(self, document: DocumentType) -> Mapping[str, int]:
    term_positions = self._index_reader.get_term_positions(document.id)
    if term_positions is None:
        raise KeyError(f"Document '{document.id}' not found in index.")
    return {term: len(positions) for term, positions in term_positions}

term_frequencies

term_frequencies(
    document: DocumentType,
) -> Mapping[str, float]
Source code in ir_axioms/tools/text_statistics/pyserini.py
def term_frequencies(self, document: DocumentType) -> Mapping[str, float]:
    term_frequencies = self._index_reader.get_document_vector(document.id)
    if term_frequencies is None:
        return {}
    return term_frequencies

TerrierDocumentTextStatistics dataclass

TerrierDocumentTextStatistics(
    *, index_location: TerrierIndex | Path | str | Any
)

Bases: TextStatistics[Document]

index_location instance-attribute

index_location: TerrierIndex | Path | str | Any

term_counts

term_counts(document: Document) -> Mapping[str, int]
Source code in ir_axioms/tools/text_statistics/pyterrier.py
def term_counts(self, document: Document) -> Mapping[str, int]:
    docid: int = self._meta_index.getDocument("docno", document.id)
    document_entry: Any = self._document_index.getDocumentEntry(docid)
    postings: Iterable[Any] | None
    postings = self._direct_index.getPostings(document_entry)
    if postings is None:
        return {}
    term_counts = {
        str(
            self._lexicon.getLexiconEntry(
                posting.getId(),
            ).getKey()
        ): posting.getFrequency()
        for posting in postings
    }
    return term_counts

SimpleTextStatistics dataclass

SimpleTextStatistics(
    *,
    text_contents: TextContents[T],
    term_tokenizer: TermTokenizer,
)

Bases: TextStatistics[T]

text_contents instance-attribute

text_contents: TextContents[T]

term_tokenizer instance-attribute

term_tokenizer: TermTokenizer

term_counts

term_counts(document: T) -> Mapping[str, int]
Source code in ir_axioms/tools/text_statistics/simple.py
def term_counts(self, document: T) -> Mapping[str, int]:
    text = self.text_contents.contents(input=document)
    if isinstance(text, TokenizedString):
        return text.tokens
    terms = self.term_tokenizer.terms_unordered(text)
    return Counter(terms)

TextStatisticsModule

Bases: Module

configure

configure(binder: Binder) -> None
Source code in ir_axioms/tools/text_statistics/__init__.py
def configure(self, binder: Binder) -> None:
    from ir_axioms.tools import TextContents, TermTokenizer

    @inject
    def _make_simple_text_statistics_query(
        text_contents: TextContents[Query],
        term_tokenizer: TermTokenizer,
    ) -> TextStatistics[Query]:
        return SimpleTextStatistics(
            text_contents=text_contents,
            term_tokenizer=term_tokenizer,
        )

    @inject
    def _make_simple_text_statistics_document(
        text_contents: TextContents[Document],
        term_tokenizer: TermTokenizer,
    ) -> TextStatistics[Document]:
        return SimpleTextStatistics(
            text_contents=text_contents,
            term_tokenizer=term_tokenizer,
        )

    @inject
    def _make_simple_text_statistics_generation_input(
        text_contents: TextContents[GenerationRequest],
        term_tokenizer: TermTokenizer,
    ) -> TextStatistics[GenerationRequest]:
        return SimpleTextStatistics(
            text_contents=text_contents,
            term_tokenizer=term_tokenizer,
        )

    @inject
    def _make_simple_text_statistics_generation_output(
        text_contents: TextContents[GenerationResponse],
        term_tokenizer: TermTokenizer,
    ) -> TextStatistics[GenerationResponse]:
        return SimpleTextStatistics(
            text_contents=text_contents,
            term_tokenizer=term_tokenizer,
        )

    binder.bind(
        interface=TextStatistics[Query],
        to=_make_simple_text_statistics_query,
        scope=singleton,
    )
    binder.bind(
        interface=TextStatistics[Document],
        to=_make_simple_text_statistics_document,
        scope=singleton,
    )
    binder.bind(
        interface=TextStatistics[GenerationRequest],
        to=_make_simple_text_statistics_generation_input,
        scope=singleton,
    )
    binder.bind(
        interface=TextStatistics[GenerationResponse],
        to=_make_simple_text_statistics_generation_output,
        scope=singleton,
    )
    binder.bind(
        interface=TextStatistics[Query | Document],
        to=DocumentQueryTextStatistics,
        scope=singleton,
    )
    binder.bind(
        interface=TextStatistics[GenerationRequest | GenerationResponse],
        to=GenerationInputOutputTextStatistics,
        scope=singleton,
    )

ir_axioms.tools.index_statistics

IndexStatistics

Bases: Protocol

document_count property

document_count: int

document_frequency

document_frequency(term: str) -> int
Source code in ir_axioms/tools/index_statistics/base.py
def document_frequency(self, term: str) -> int: ...

document_frequencies

document_frequencies(
    terms: Iterable[str],
) -> Mapping[str, int]
Source code in ir_axioms/tools/index_statistics/base.py
def document_frequencies(self, terms: Iterable[str]) -> Mapping[str, int]:
    return {term: self.document_frequency(term) for term in terms}

inverse_document_frequency

inverse_document_frequency(term: str) -> float
Source code in ir_axioms/tools/index_statistics/base.py
def inverse_document_frequency(self, term: str) -> float:
    return self._inverse_document_frequency(self.document_frequency(term))

inverse_document_frequencies

inverse_document_frequencies(
    terms: Iterable[str],
) -> Mapping[str, float]
Source code in ir_axioms/tools/index_statistics/base.py
def inverse_document_frequencies(self, terms: Iterable[str]) -> Mapping[str, float]:
    return {
        term: self._inverse_document_frequency(document_frequency)
        for term, document_frequency in self.document_frequencies(terms).items()
    }

AnseriniIndexStatistics dataclass

AnseriniIndexStatistics(
    *,
    index_dir: Path | str,
    analyzer: Analyzer = default_analyzer(),
)

Bases: IndexStatistics

index_dir instance-attribute

index_dir: Path | str

analyzer class-attribute instance-attribute

analyzer: Analyzer = field(default_factory=default_analyzer)

document_count cached property

document_count: int

document_frequency

document_frequency(term: str) -> int
Source code in ir_axioms/tools/index_statistics/pyserini.py
def document_frequency(self, term: str) -> int:
    if term not in self._document_frequencies:
        document_frequency, _ = self._index_reader.get_term_counts(
            term=term,
            analyzer=self.analyzer.analyzer,
        )
        self._document_frequencies[term] = document_frequency
    return self._document_frequencies[term]

TerrierIndexStatistics dataclass

TerrierIndexStatistics(
    *, index_location: TerrierIndex | Path | str | Any
)

Bases: IndexStatistics

index_location instance-attribute

index_location: TerrierIndex | Path | str | Any

document_count cached property

document_count: int

document_frequency

document_frequency(term: str) -> int
Source code in ir_axioms/tools/index_statistics/pyterrier.py
def document_frequency(self, term: str) -> int:
    if term not in self._document_frequencies:
        entry = self._lexicon.getLexiconEntry(term)
        if entry is None or entry.getNumberOfEntries() == 0:
            document_frequency = 0
        else:
            document_frequency = int(entry.getDocumentFrequency())
        self._document_frequencies[term] = document_frequency
    return self._document_frequencies[term]

ir_axioms.tools.similarity

TermSimilarity

Bases: Protocol

similarity

similarity(term1: str, term2: str) -> float
Source code in ir_axioms/tools/similarity/base.py
def similarity(self, term1: str, term2: str) -> float: ...

similarities

similarities(terms: Sequence[str]) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/base.py
def similarities(self, terms: Sequence[str]) -> NDArray[float_]:
    return array(
        [
            self.similarity(
                term1=term1,
                term2=term2,
            )
            for term1 in terms
            for term2 in terms
        ],
        dtype=float_,
    ).reshape((len(terms), len(terms)))

similarity_sums

similarity_sums(terms: Iterable[str]) -> dict[str, float]
Source code in ir_axioms/tools/similarity/base.py
def similarity_sums(self, terms: Iterable[str]) -> dict[str, float]:
    similarity_sums: dict[str, float] = {term: 0 for term in terms}
    for term1, term2 in combinations(similarity_sums.keys(), 2):
        similarity = self.similarity(term1, term2)
        similarity_sums[term1] += similarity
        similarity_sums[term2] += similarity
    return similarity_sums

average_similarity

average_similarity(
    terms1: Collection[str], terms2: Collection[str]
) -> float
Source code in ir_axioms/tools/similarity/base.py
def average_similarity(
    self, terms1: Collection[str], terms2: Collection[str]
) -> float:
    if len(terms1) == 0 or len(terms2) == 0:
        return 0
    return array(
        [self.similarity(term1, term2) for term1 in terms1 for term2 in terms2]
    ).mean()

max_similarity_pairs

max_similarity_pairs(
    terms1: Collection[str], terms2: Collection[str]
) -> tuple[float, Set[tuple[str, str]]]

Most similar pairs of terms from the two collections, along with their similarity (or -inf if there are no pairs).

Source code in ir_axioms/tools/similarity/base.py
def max_similarity_pairs(
    self, terms1: Collection[str], terms2: Collection[str]
) -> tuple[float, Set[tuple[str, str]]]:
    """
    Most similar pairs of terms from the two collections,
    along with their similarity (or -inf if there are no pairs).
    """
    if len(terms1) == 0 or len(terms2) == 0:
        return -inf, set()
    most_similar_pairs: Sequence[tuple[str, str]] = sorted(
        product(set(terms1), set(terms2)),
        key=self._pair_similarity,
        reverse=True,
    )
    max_similarity = self._pair_similarity(most_similar_pairs[0])
    return max_similarity, {
        pair
        for pair in most_similar_pairs
        if self._pair_similarity(pair) >= max_similarity
    }

max_average_similarity_terms

max_average_similarity_terms(
    terms: Collection[str],
) -> Set[str]
Source code in ir_axioms/tools/similarity/base.py
def max_average_similarity_terms(self, terms: Collection[str]) -> Set[str]:
    if len(terms) == 0:
        return set()
    similarity_sums = self.similarity_sums(set(terms))
    max_similarity_sum = max(similarity_sums.values())
    return {term for term in terms if similarity_sums[term] >= max_similarity_sum}

min_average_similarity_terms

min_average_similarity_terms(
    terms: Collection[str],
) -> Set[str]
Source code in ir_axioms/tools/similarity/base.py
def min_average_similarity_terms(self, terms: Collection[str]) -> Set[str]:
    if len(terms) == 0:
        return set()
    similarity_sums = self.similarity_sums(set(terms))
    min_similarity_sum = min(similarity_sums.values())
    return {term for term in terms if similarity_sums[term] <= min_similarity_sum}

SentenceSimilarity

Bases: Protocol

similarity

similarity(sentence1: str, sentence2: str) -> float
Source code in ir_axioms/tools/similarity/base.py
def similarity(self, sentence1: str, sentence2: str) -> float: ...

similarities

similarities(sentences: Sequence[str]) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/base.py
def similarities(self, sentences: Sequence[str]) -> NDArray[float_]:
    return array(
        [
            self.similarity(
                sentence1=sentence1,
                sentence2=sentence2,
            )
            for sentence1 in sentences
            for sentence2 in sentences
        ],
        dtype=float_,
    ).reshape((len(sentences), len(sentences)))

self_similarities

self_similarities(
    sentences: Sequence[str],
) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/base.py
@deprecated("Use similarities instead.")
def self_similarities(self, sentences: Sequence[str]) -> NDArray[float_]:
    return self.similarities(sentences)

paired_similarities

paired_similarities(
    sentences1: Sequence[str], sentences2: Sequence[str]
) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/base.py
def paired_similarities(
    self, sentences1: Sequence[str], sentences2: Sequence[str]
) -> NDArray[float_]:
    return array(
        [
            self.similarity(
                sentence1=sentence1,
                sentence2=sentence2,
            )
            for sentence1 in sentences1
            for sentence2 in sentences2
        ],
        dtype=float_,
    ).reshape((len(sentences1), len(sentences2)))

average_similarity

average_similarity(
    sentences1: Collection[str], sentences2: Collection[str]
) -> float
Source code in ir_axioms/tools/similarity/base.py
def average_similarity(
    self, sentences1: Collection[str], sentences2: Collection[str]
) -> float:
    return self.paired_similarities(list(sentences1), list(sentences2)).mean()

FastTextTermSimilarity dataclass

FastTextTermSimilarity(
    model_name: str = "facebook/fasttext-en-vectors",
)

Bases: TermSimilarity

model_name class-attribute instance-attribute

model_name: str = 'facebook/fasttext-en-vectors'

model cached property

model: _FastText

similarity

similarity(term1: str, term2: str) -> float
Source code in ir_axioms/tools/similarity/fasttext.py
@final
def similarity(self, term1: str, term2: str) -> float:
    key = (term1, term2)
    if key not in self._similarities:
        vector1 = self.model.get_word_vector(term1)
        vector2 = self.model.get_word_vector(term2)
        self._similarities[key] = _cosine_similarity(vector1, vector2)
    return self._similarities[key]

WordNetSynonymSetTermSimilarity dataclass

WordNetSynonymSetTermSimilarity(smoothing: int = 0)

Bases: TermSimilarity

smoothing class-attribute instance-attribute

smoothing: int = 0

similarity

similarity(term1: str, term2: str) -> float
Source code in ir_axioms/tools/similarity/wordnet.py
@final
def similarity(self, term1: str, term2: str) -> float:
    return _synonym_set_similarity(term1, term2, self.smoothing)

AverageTermSimilaritySentenceSimilarity dataclass

AverageTermSimilaritySentenceSimilarity(
    term_tokenizer: TermTokenizer,
    term_similarity: TermSimilarity,
)

Bases: SentenceSimilarity

Sentence similarity based on the mean term similarity of the terms in the sentences.

Note: This is only meant as a fallback if no other sentence similarity model is available.

term_tokenizer instance-attribute

term_tokenizer: TermTokenizer

term_similarity instance-attribute

term_similarity: TermSimilarity

similarity

similarity(sentence1: str, sentence2: str) -> float
Source code in ir_axioms/tools/similarity/simple.py
def similarity(self, sentence1: str, sentence2: str) -> float:
    terms1 = self.term_tokenizer.terms_unordered(sentence1)
    terms2 = self.term_tokenizer.terms_unordered(sentence2)
    return self.term_similarity.average_similarity(terms1, terms2)

Model2VecSentenceSimilarity dataclass

Model2VecSentenceSimilarity(
    model_name: str = "minishlab/potion-base-32M",
)

Bases: SentenceSimilarity

model_name class-attribute instance-attribute

model_name: str = 'minishlab/potion-base-32M'

model cached property

model: StaticModel

similarity

similarity(sentence1: str, sentence2: str) -> float
Source code in ir_axioms/tools/similarity/model2vec.py
def similarity(self, sentence1: str, sentence2: str) -> float:
    return self.similarities([sentence1, sentence2])[0, 1]

similarities

similarities(sentences: Sequence[str]) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/model2vec.py
def similarities(self, sentences: Sequence[str]) -> NDArray[float_]:
    vectors = self.model.encode(
        sentences=list(sentences),
        show_progress_bar=False,
    )
    return array(
        [
            _cosine_similarity(vectors[i1], vectors[i2])
            for i1 in range(len(sentences))
            for i2 in range(len(sentences))
        ],
        dtype=float_,
    ).reshape((len(sentences), len(sentences)))

paired_similarities

paired_similarities(
    sentences1: Sequence[str], sentences2: Sequence[str]
) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/model2vec.py
def paired_similarities(
    self, sentences1: Sequence[str], sentences2: Sequence[str]
) -> NDArray[float_]:
    vectors1 = self.model.encode(
        sentences=list(sentences1),
        show_progress_bar=False,
    )
    vectors2 = self.model.encode(
        sentences=list(sentences2),
        show_progress_bar=False,
    )
    return array(
        [
            _cosine_similarity(vector1, vector2)
            for vector1 in vectors1
            for vector2 in vectors2
        ],
        dtype=float_,
    ).reshape((len(sentences1), len(sentences2)))

SentenceTransformersSentenceSimilarity dataclass

SentenceTransformersSentenceSimilarity(
    model_name: str = "sentence-transformers/all-mpnet-base-v2",
)

Bases: SentenceSimilarity

model_name class-attribute instance-attribute

model_name: str = 'sentence-transformers/all-mpnet-base-v2'

model cached property

model: SentenceTransformer

similarity

similarity(sentence1: str, sentence2: str) -> float
Source code in ir_axioms/tools/similarity/sentence_transformers.py
def similarity(self, sentence1: str, sentence2: str) -> float:
    return self.similarities([sentence1, sentence2])[0, 1]

similarities

similarities(sentences: Sequence[str]) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/sentence_transformers.py
def similarities(self, sentences: Sequence[str]) -> NDArray[float_]:
    vectors = self.model.encode(
        inputs=list(sentences),
        convert_to_numpy=True,
    )
    return array(
        [
            _cosine_similarity(vectors[i1], vectors[i2])
            for i1 in range(len(sentences))
            for i2 in range(len(sentences))
        ],
        dtype=float_,
    ).reshape((len(sentences), len(sentences)))

paired_similarities

paired_similarities(
    sentences1: Sequence[str], sentences2: Sequence[str]
) -> NDArray[float_]
Source code in ir_axioms/tools/similarity/sentence_transformers.py
def paired_similarities(
    self, sentences1: Sequence[str], sentences2: Sequence[str]
) -> NDArray[float_]:
    vectors1 = self.model.encode(
        inputs=list(sentences1),
        convert_to_numpy=True,
    )
    vectors2 = self.model.encode(
        inputs=list(sentences2),
        convert_to_numpy=True,
    )
    return array(
        [
            _cosine_similarity(vector1, vector2)
            for vector1 in vectors1
            for vector2 in vectors2
        ],
        dtype=float_,
    ).reshape((len(sentences1), len(sentences2)))

SimilarityModule

Bases: Module

configure

configure(binder: Binder) -> None
Source code in ir_axioms/tools/similarity/__init__.py
def configure(self, binder: Binder) -> None:
    binder.bind(
        interface=TermSimilarity,
        to=WordNetSynonymSetTermSimilarity,
        scope=singleton,
    )
    if is_fasttext_installed():
        binder.bind(
            interface=TermSimilarity,
            to=FastTextTermSimilarity,
            scope=singleton,
        )

    binder.bind(
        interface=SentenceSimilarity,
        to=AverageTermSimilaritySentenceSimilarity,
        scope=singleton,
    )
    if is_model2vec_installed():
        binder.bind(
            interface=SentenceSimilarity,
            to=Model2VecSentenceSimilarity,
            scope=singleton,
        )
    if is_sentence_transformers_installed():
        binder.bind(
            interface=SentenceSimilarity,
            to=SentenceTransformersSentenceSimilarity,
            scope=singleton,
        )

ir_axioms.tools.aspects

AspectExtraction

Bases: Protocol

aspects

aspects(text: str) -> Set[str]
Source code in ir_axioms/tools/aspects/base.py
def aspects(self, text: str) -> Set[str]: ...

iter_aspects

iter_aspects(texts: Iterable[str]) -> Iterator[Set[str]]
Source code in ir_axioms/tools/aspects/base.py
def iter_aspects(self, texts: Iterable[str]) -> Iterator[Set[str]]:
    for text in texts:
        yield self.aspects(text)

KeyBertAspectExtraction dataclass

KeyBertAspectExtraction(
    model: str = "all-MiniLM-L6-v2",
    ngram_range: tuple[int, int] = (1, 3),
    top_n: int = 10,
    threshold: float = 0.0,
)

Bases: AspectExtraction

model class-attribute instance-attribute

model: str = 'all-MiniLM-L6-v2'

ngram_range class-attribute instance-attribute

ngram_range: tuple[int, int] = (1, 3)

top_n class-attribute instance-attribute

top_n: int = 10

threshold class-attribute instance-attribute

threshold: float = 0.0

aspects

aspects(text: str) -> Set[str]
Source code in ir_axioms/tools/aspects/keybert.py
def aspects(self, text: str) -> Set[str]:
    keyphrases: list[tuple[str, float]] = self._model.extract_keywords(
        docs=text,
        keyphrase_ngram_range=self.ngram_range,
        top_n=self.top_n,
    )  # type: ignore
    return {phrase for phrase, score in keyphrases if score >= self.threshold}

iter_aspects

iter_aspects(texts: Iterable[str]) -> Iterator[Set[str]]
Source code in ir_axioms/tools/aspects/keybert.py
def iter_aspects(self, texts: Iterable[str]) -> Iterator[Set[str]]:
    keyphrases: list[list[tuple[str, float]]] = self._model.extract_keywords(
        docs=list(texts),
        keyphrase_ngram_range=self.ngram_range,
        top_n=self.top_n,
    )  # type: ignore
    for keys in keyphrases:
        yield {phrase for phrase, score in keys if score >= self.threshold}

SpacyNounChunksAspectExtraction dataclass

SpacyNounChunksAspectExtraction(
    language_name: str = "en_core_web_sm",
    chunks: Literal["direct", "root", "both"] = "direct",
)

Bases: AspectExtraction

language_name class-attribute instance-attribute

language_name: str = 'en_core_web_sm'

chunks class-attribute instance-attribute

chunks: Literal['direct', 'root', 'both'] = 'direct'

aspects

aspects(text: str) -> Set[str]
Source code in ir_axioms/tools/aspects/spacy.py
def aspects(self, text: str) -> Set[str]:
    document = self._language(text)
    chunks = []
    if self.chunks == "direct" or self.chunks == "both":
        chunks += [chunk.text.lower() for chunk in document.noun_chunks]
    if self.chunks == "root" or self.chunks == "both":
        chunks += [chunk.root.text.lower() for chunk in document.noun_chunks]
    return set(chunks)

SpacyEntitiesAspectExtraction dataclass

SpacyEntitiesAspectExtraction(
    language_name: str = "en_core_web_sm",
    chunks: Literal["direct", "root", "both"] = "direct",
)

Bases: AspectExtraction

language_name class-attribute instance-attribute

language_name: str = 'en_core_web_sm'

chunks class-attribute instance-attribute

chunks: Literal['direct', 'root', 'both'] = 'direct'

aspects

aspects(text: str) -> Set[str]
Source code in ir_axioms/tools/aspects/spacy.py
def aspects(self, text: str) -> Set[str]:
    document = self._language(text)
    return {entity.text for entity in document.ents}

iter_aspects

iter_aspects(texts: Iterable[str]) -> Iterator[Set[str]]
Source code in ir_axioms/tools/aspects/spacy.py
def iter_aspects(self, texts: Iterable[str]) -> Iterator[Set[str]]:
    documents = self._language.pipe(texts)
    for document in documents:
        yield {entity.text for entity in document.ents}

YakeAspectExtraction dataclass

YakeAspectExtraction(
    language_name: str = "en_core_web_sm",
    ngram_range: tuple[int, int] = (1, 3),
    include_pos: Iterable[str] = ("NOUN", "PROPN", "ADJ"),
    top_n: int = 10,
    threshold: float = inf,
)

Bases: AspectExtraction

Extract key phrases using YAKE (via textacy) as aspects.

For texts with fewer than two distinct content words (e.g., single-word queries), YAKE cannot score candidates. Then, the (lemmatized) unigram candidates are returned without scoring.

language_name class-attribute instance-attribute

language_name: str = 'en_core_web_sm'

ngram_range class-attribute instance-attribute

ngram_range: tuple[int, int] = (1, 3)

include_pos class-attribute instance-attribute

include_pos: Iterable[str] = ('NOUN', 'PROPN', 'ADJ')

top_n class-attribute instance-attribute

top_n: int = 10

threshold class-attribute instance-attribute

threshold: float = inf

aspects

aspects(text: str) -> Set[str]
Source code in ir_axioms/tools/aspects/textacy.py
def aspects(self, text: str) -> Set[str]:
    return self._extract(self._language(text))

iter_aspects

iter_aspects(texts: Iterable[str]) -> Iterator[Set[str]]
Source code in ir_axioms/tools/aspects/textacy.py
def iter_aspects(self, texts: Iterable[str]) -> Iterator[Set[str]]:
    for document in self._language.pipe(texts):
        yield self._extract(document)

AspectsModule

Bases: Module

configure

configure(binder: Binder) -> None
Source code in ir_axioms/tools/aspects/__init__.py
def configure(self, binder: Binder) -> None:
    binder.bind(
        interface=AspectExtraction,
        to=SpacyNounChunksAspectExtraction,
        scope=singleton,
    )

ir_axioms.tools.pivot

PivotSelection

Bases: Protocol, Generic[Input, Output]

select_pivot

select_pivot(
    input: Input, vertices: Sequence[Output]
) -> Output
Source code in ir_axioms/tools/pivot/base.py
def select_pivot(self, input: Input, vertices: Sequence[Output]) -> Output: ...

RandomPivotSelection dataclass

RandomPivotSelection(
    *,
    seed: int
    | float
    | str
    | bytes
    | bytearray
    | None = None,
)

Bases: PivotSelection[Input, Output]

seed class-attribute instance-attribute

seed: int | float | str | bytes | bytearray | None = None

select_pivot

select_pivot(
    input: Input, vertices: Sequence[Output]
) -> Output
Source code in ir_axioms/tools/pivot/simple.py
def select_pivot(self, input: Input, vertices: Sequence[Output]) -> Output:
    return vertices[self._random.randint(0, len(vertices) - 1)]

FirstPivotSelection dataclass

FirstPivotSelection()

Bases: PivotSelection[Input, Output]

select_pivot

select_pivot(
    input: Input, vertices: Sequence[Output]
) -> Output
Source code in ir_axioms/tools/pivot/simple.py
def select_pivot(self, input: Input, vertices: Sequence[Output]) -> Output:
    return vertices[0]

LastPivotSelection dataclass

LastPivotSelection()

Bases: PivotSelection[Input, Output]

select_pivot

select_pivot(
    input: Input, vertices: Sequence[Output]
) -> Output
Source code in ir_axioms/tools/pivot/simple.py
def select_pivot(self, input: Input, vertices: Sequence[Output]) -> Output:
    return vertices[-1]

MiddlePivotSelection dataclass

MiddlePivotSelection()

Bases: PivotSelection[Input, Output]

select_pivot

select_pivot(
    input: Input, vertices: Sequence[Output]
) -> Output
Source code in ir_axioms/tools/pivot/simple.py
def select_pivot(self, input: Input, vertices: Sequence[Output]) -> Output:
    return vertices[(len(vertices) - 1) // 2]

PivotModule

Bases: Module

configure

configure(binder: Binder) -> None
Source code in ir_axioms/tools/pivot/__init__.py
def configure(self, binder: Binder) -> None:
    binder.bind(
        interface=PivotSelection,
        to=RandomPivotSelection,
        scope=singleton,
    )