Public/kotaemon
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: tei_fast_rerank logic bump:patch

Browse Source
This commit is contained in:
Tadashi 2024-12-17 17:33:47 +07:00
parent 54320d08df
commit 3bd19f399f
No known key found for this signature in database
GPG Key ID: 399380A00CC9028D
1 changed files with 10 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

14
libs/kotaemon/kotaemon/rerankings/tei_fast_rerank.py
View File

@ -29,12 +29,18 @@ class TeiFastReranking(BaseReranking):
), ),
) )
is_truncated: Optional[bool] = Param(True, help="Whether to truncate the inputs") is_truncated: Optional[bool] = Param(True, help="Whether to truncate the inputs")
max_tokens: Optional[int] = Param(512, help="This option is used to specify the maximum number of tokens supported by the reranker model.") max_tokens: Optional[int] = Param(
512,
help=(
"This option is used to specify the "
"maximum number of tokens supported by the reranker model."
),
)
def client(self, query, texts): def client(self, query, texts):
if self.is_truncated == True: if self.is_truncated:
max_tokens = self.max_tokens # default is 512 tokens. max_tokens = self.max_tokens # default is 512 tokens.
truncated_texts = [text[:max_tokens] for text in texts] truncated_texts = [text[:max_tokens] for text in texts]
response = session.post( response = session.post(
url=self.endpoint_url, url=self.endpoint_url,