Source code for lightning_ir.models.t5.config
1from typing import Literal
2
3from ...cross_encoder.config import CrossEncoderConfig
4
5
[docs]
6class T5CrossEncoderConfig(CrossEncoderConfig):
7
8 model_type = "encoder-decoder-cross-encoder"
9
10 TOKENIZER_ARGS = CrossEncoderConfig.TOKENIZER_ARGS.union({"decoder_strategy"})
11 ADDED_ARGS = CrossEncoderConfig.ADDED_ARGS.union(TOKENIZER_ARGS)
12
[docs]
13 def __init__(
14 self,
15 query_length: int = 32,
16 doc_length: int = 512,
17 decoder_strategy: Literal["mono", "rank"] = "mono",
18 **kwargs,
19 ) -> None:
20 kwargs["pooling_strategy"] = "first"
21 super().__init__(query_length=query_length, doc_length=doc_length, **kwargs)
22 self.decoder_strategy = decoder_strategy
