Skip to content

Commit 3bd19f3

Browse files
committed
fix: tei_fast_rerank logic bump:patch
1 parent 54320d0 commit 3bd19f3

File tree

1 file changed

+10
-4
lines changed

1 file changed

+10
-4
lines changed

libs/kotaemon/kotaemon/rerankings/tei_fast_rerank.py

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -29,12 +29,18 @@ class TeiFastReranking(BaseReranking):
2929
),
3030
)
3131
is_truncated: Optional[bool] = Param(True, help="Whether to truncate the inputs")
32-
max_tokens: Optional[int] = Param(512, help="This option is used to specify the maximum number of tokens supported by the reranker model.")
32+
max_tokens: Optional[int] = Param(
33+
512,
34+
help=(
35+
"This option is used to specify the "
36+
"maximum number of tokens supported by the reranker model."
37+
),
38+
)
3339

3440
def client(self, query, texts):
35-
if self.is_truncated == True:
36-
max_tokens = self.max_tokens # default is 512 tokens.
37-
truncated_texts = [text[:max_tokens] for text in texts]
41+
if self.is_truncated:
42+
max_tokens = self.max_tokens # default is 512 tokens.
43+
truncated_texts = [text[:max_tokens] for text in texts]
3844

3945
response = session.post(
4046
url=self.endpoint_url,

0 commit comments

Comments
 (0)