初始化项目，由ModelHub XC社区提供模型

Model: iic/ERank-4B
Source: Original Platform

examples/ERank_Transformer.py

@@ -0,0 +1,146 @@
+from torch.nn import functional as F
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from utils import prompt_template, truncate, hybrid_scores
+
+class ERank_Transformer:
+    
+    def __init__(self, model_name_or_path: str):
+        """
+        Initializes the ERank_Transformer reranker.
+
+        Args:
+            model_name_or_path (str): The name or path of the model to be loaded.
+                                      This can be a Hugging Face model ID or a local path.
+        """
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+        self.reranker = AutoModelForCausalLM.from_pretrained(model_name_or_path).eval()
+        self.reranker.to("cuda")
+  
+    def rerank(self, query: str, docs: list, instruction: str, truncate_length: int=None) -> list:
+        """
+        Reranks a list of documents based on a query and a specific instruction.
+
+        Args:
+            query (str): The search query provided by the user.
+            docs (list): A list of dictionaries, where each dictionary represents a document
+                         and must contain a "content" key.
+            instruction (str): The instruction for the model, guiding it on how to evaluate the documents.
+            truncate_length (int, optional): The maximum length to truncate the query and document content to. Defaults to None.
+
+        Returns:
+            list: A new list of document dictionaries, sorted by their "rank_score" in descending order.
+        """
+        
+        # prepare messages
+        messages = [
+            [{
+                "role": "user", 
+                "content": prompt_template.format(
+                    query=truncate(self.tokenizer, query, length=truncate_length) if truncate_length else query,
+                    doc=truncate(self.tokenizer, doc["content"], length=truncate_length) if truncate_length else doc["content"],
+                    instruction=instruction
+                )
+            }] for doc in docs
+        ]
+
+        # encode tokens
+        texts = [
+            self.tokenizer.apply_chat_template(
+                each,
+                tokenize=False,
+                add_generation_prompt=True,
+            ) for each in messages
+        ]
+        inputs = self.tokenizer(texts, padding=True, return_tensors="pt").to(self.reranker.device)
+
+        # LLM completion
+        outputs = self.reranker.generate(
+            **inputs,
+            max_new_tokens=8192,
+            output_scores=True,
+            return_dict_in_generate=True
+        )
+  
+        # extract and organize results
+        results = []
+        scores = outputs.scores
+        generated_ids = outputs.sequences
+        answer_token_ids = self.tokenizer.encode("<answer>", add_special_tokens=False)
+        for idx in range(len(texts)):
+            
+            # find <answer> in the generated sequence
+            output_ids = generated_ids[idx].tolist()
+            start_index = -1
+            for i in range(len(output_ids)-len(answer_token_ids)-1, -1, -1):
+                if output_ids[i:i + len(answer_token_ids)] == answer_token_ids:
+                    start_index = i + len(answer_token_ids)
+                    break
+            
+            # start from the index after <answer>
+            answer = ""
+            prob = 1.0
+            if start_index != -1:
+                for t in range(start_index - inputs.input_ids.size(1), len(scores)):
+                    generated_token_id = generated_ids[idx][inputs.input_ids.size(1) + t]
+                    token = self.tokenizer.decode(generated_token_id)
+                    if token.isdigit():
+                        logits = scores[t][idx]
+                        probs = F.softmax(logits, dim=-1)
+                        prob *= probs[generated_token_id].item()
+                        answer += token
+                    else:
+                        break
+
+            # in case the answer is not a digit or exceeds 10
+            try:
+                answer = int(answer)
+                assert answer <= 10
+            except:
+                answer = -1
+                
+            # append to the final results
+            results.append({
+                **docs[idx],
+                "rank_score": answer * prob
+            })
+            
+        # sort the reranking results for the query
+        results.sort(key=lambda x:x["rank_score"], reverse=True)
+        return results
+    
+    
+if __name__ == "__main__":
+    
+    # select a model
+    model_name_or_path = "Ucreate/ERank-4B"
+    # model_name_or_path = "Ucreate/ERank-14B"
+    # model_name_or_path = "Ucreate/ERank-32B"
+    reranker = ERank_Transformer(model_name_or_path)
+    
+    # input data
+    instruction = "Retrieve relevant documents for the query."
+    query = "I am happy"
+    docs = [
+        {"content": "excited", "first_stage_score": 46.7},
+        {"content": "sad", "first_stage_score": 1.5},
+        {"content": "peaceful", "first_stage_score": 2.3},
+    ]
+
+    # rerank
+    results = reranker.rerank(query, docs, instruction, truncate_length=2048)
+    print(results)
+    # [
+    #	 {'content': 'excited', 'first_stage_score': 46.7, 'rank_score': 4.84}, 
+    #	 {'content': 'peaceful', 'first_stage_score': 2.3, 'rank_score': 2.98}
+    #	 {'content': 'sad', 'first_stage_score': 1.5, 'rank_score': 0.0}, 
+    # ]
+    
+    # Optional: hybrid with first-stage scores
+    alpha = 0.2
+    hybrid_results = hybrid_scores(results, alpha)
+    print(hybrid_results)
+    # [
+    #	 {'content': 'excited', 'first_stage_score': 46.7, 'rank_score': 4.84, 'hybrid_score': 1.18}, 
+    #	 {'content': 'peaceful', 'first_stage_score': 2.3, 'rank_score': 2.98, 'hybrid_score':0.01}, 
+    #	 {'content': 'sad', 'first_stage_score': 1.5, 'rank_score': 0.0, 'hybrid_score': -1.19}
+    # ]