fix: tei_fast_rerank logic bump:patch

Cinnamon · Dec 17, 2024 · 3bd19f3 · 3bd19f3
1 parent 54320d0
commit 3bd19f3
Showing 1 changed file with 10 additions and 4 deletions.
diff --git a/libs/kotaemon/kotaemon/rerankings/tei_fast_rerank.py b/libs/kotaemon/kotaemon/rerankings/tei_fast_rerank.py
@@ -29,12 +29,18 @@ class TeiFastReranking(BaseReranking):
         ),
     )
     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):
-        if self.is_truncated == True:
-                max_tokens = self.max_tokens # default is 512 tokens.
-                truncated_texts = [text[:max_tokens] for text in texts]
+        if self.is_truncated:
+            max_tokens = self.max_tokens  # default is 512 tokens.
+            truncated_texts = [text[:max_tokens] for text in texts]
 
         response = session.post(
             url=self.endpoint_url,