minor change

reranker.py

@@ -14,7 +14,14 @@ class Reranker:
         return output
       
     def process_inputs(self,pairs):
-        prefix = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n"
+        prefix = (
+            "<|im_start|>system\n"
+            "Procijeni da li dati Dokument adekvatno odgovara na Upit na osnovu pravne instrukcije. "
+            "Dokument treba da bude relevantan, tačan i u skladu sa važećim pravnim propisima i standardima. "
+            "Odgovor mora biti striktno \"da\" ako ispunjava uslove, ili \"ne\" ako ne ispunjava.\n"
+            "<|im_end|>\n"
+            "<|im_start|>user\n"
+        )        
         suffix = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n"
         prefix_tokens = self.tokenizer.encode(prefix, add_special_tokens=False)
         suffix_tokens = self.tokenizer.encode(suffix, add_special_tokens=False)
@@ -35,7 +42,7 @@ class Reranker:
         token_false_id = self.tokenizer.convert_tokens_to_ids("no")
         token_true_id = self.tokenizer.convert_tokens_to_ids("yes")
 
-        task = 'Given a web search query, retrieve relevant passages that answer the query'
+        task = 'Na osnovu datog upita, vrati najrelevantije rezultate koje odgovaraju upitu'
         pairs = [self.format_instruction(task, query, doc) for query, doc in zip(queries, documents)]
         inputs = self.process_inputs(pairs)
         
@@ -46,4 +53,8 @@ class Reranker:
         batch_scores = torch.nn.functional.log_softmax(batch_scores, dim=1)
         scores = batch_scores[:, 1].exp().tolist()
 
-        return scores
+        results = list(zip(scores, queries, documents))
+        results.sort(key=lambda x: x[0], reverse=True)
+        top_10 = results[:10]
+
+        return top_10