restapi

app.py

@@ -173,26 +173,22 @@ class QwenReranker(RerankerModel):
         )
     
     def _process_inputs(self, pairs: List[str]):
-        """Process input pairs for Qwen model."""
+        """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
+        ]
+
         inputs = self.tokenizer(
-            pairs,
-            padding="max_length",        # ✅ pad langsung ke max_length
-            truncation=True,             # ✅ aktifkan truncation
-            max_length=self.max_length - len(self.prefix_tokens) - len(self.suffix_tokens),
+            modified_pairs,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_length,
             return_tensors="pt"
-        )
-
-        # Tambahkan prefix & suffix tokens
-        for i, ele in enumerate(inputs['input_ids']):
-            inputs['input_ids'][i] = torch.tensor(
-                self.prefix_tokens + ele.tolist() + self.suffix_tokens, 
-                dtype=torch.long
-            )
+        ).to(self.model.device)
 
-        inputs = {k: v.to(self.model.device) for k, v in inputs.items()}
-
-        
         return inputs
+
     
     @torch.no_grad()
     def _compute_logits(self, inputs):