Source code for lightning_ir.models.xtr.config

 1from typing import Literal
 2
 3from ..col import ColConfig
 4
 5


[docs]
 6class XTRConfig(ColConfig):
 7    model_type = "xtr"
 8
 9    ADDED_ARGS = ColConfig.ADDED_ARGS.union({"token_retrieval_k", "fill_strategy", "normalization"})
10


[docs]
11    def __init__(
12        self,
13        token_retrieval_k: int | None = None,
14        fill_strategy: Literal["zero", "min"] = "zero",
15        normalization: Literal["Z"] | None = "Z",
16        **kwargs
17    ) -> None:
18        super().__init__(**kwargs)
19        self.token_retrieval_k = token_retrieval_k
20        self.fill_strategy = fill_strategy
21        self.normalization = normalization