restapi

app.py

@@ -173,22 +173,29 @@ class QwenReranker(RerankerModel):
         )
     
     def _process_inputs(self, pairs: List[str]):
-        """Process input pairs for Qwen model with prefix and suffix."""
-        modified_pairs = [
-            self.tokenizer.decode(self.prefix_tokens) + text + self.tokenizer.decode(self.suffix_tokens)
-            for text in pairs
-        ]
-
+        """Process input pairs for Qwen model."""
         inputs = self.tokenizer(
-            modified_pairs,
-            padding="max_length",
-            truncation=True,
-            max_length=self.max_length,
-            return_tensors="pt"
-        ).to(self.model.device)
-
+            pairs, 
+            padding=False, 
+            truncation='longest_first',
+            return_attention_mask=False, 
+            max_length=self.max_length - len(self.prefix_tokens) - len(self.suffix_tokens)
+        )
+        
+        for i, ele in enumerate(inputs['input_ids']):
+            inputs['input_ids'][i] = self.prefix_tokens + ele + self.suffix_tokens
+        
+        inputs = self.tokenizer.pad(
+            inputs, 
+            padding=True, 
+            return_tensors="pt", 
+            max_length=self.max_length
+        )
+        
+        for key in inputs:
+            inputs[key] = inputs[key].to(self.model.device)
+        
         return inputs
-
     
     @torch.no_grad()
     def _compute_logits(self, inputs):