Create RM_demo.py

RM_demo.py

@@ -0,0 +1,54 @@
+import logging
+from typing import List
+
+import torch
+import torch.nn as nn
+from transformers import AutoTokenizer, AutoConfig, MistralForCausalLM
+from safetensors.torch import load_file
+
+class RewardModel:
+    def __init__(self, model_dir) -> None:
+        config = AutoConfig.from_pretrained(model_dir)
+        # config._attn_implementation = "flash_attention_2"
+        self.device = torch.device('cuda')
+        self.model = MistralForCausalLM(config)
+        self.model.lm_head = nn.Linear(config.hidden_size, 1, bias=False)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_dir)
+        state_dict = load_file(f"{model_dir}/model.safetensors")
+        self.model.load_state_dict(state_dict, strict=False)
+        self.model.to(dtype=torch.bfloat16)
+        self.model.to(device=self.device)
+        self.model.eval()
+        logging.info("Load model completed.")
+
+    @torch.no_grad()
+    def score(self, prompts, chosens) -> List[float]:
+        # Concat prompt and chosen, append eos_id
+        input_ids_list = [self.tokenizer.encode(prompt) + self.tokenizer.encode(chosen) + [self.tokenizer.eos_token_id] for prompt, chosen in zip(prompts, chosens)]
+
+        # Pad sequences to the maximum length
+        max_length = max(len(ids) for ids in input_ids_list)
+        padded_input_ids = [ids + [self.tokenizer.pad_token_id or self.tokenizer.eos_token_id] * (max_length - len(ids)) for ids in input_ids_list]
+
+        # Forward pass
+        input_ids = torch.tensor(padded_input_ids).to(device=self.device)
+        logits = self.model(input_ids).logits
+
+        # Extract logits corresponding to eos_token_id positions
+        scores = []
+        for i, input_ids in enumerate(input_ids_list):
+            eos_position = input_ids.index(self.tokenizer.eos_token_id)
+            eos_logit = logits[i, eos_position, :].squeeze().item()
+            scores.append(eos_logit)
+
+        return scores
+
+if __name__ == '__main__':
+
+    local_model_dir = "Your local model dir"
+    model_dir = f"{local_model_dir}/Seed-X-RM-7B"  
+    prompt = ["Translate the following English sentence into Chinese:\nMay the force be with you <zh>"]
+    candidate = ["愿原力与你同在"]
+    model = RewardModel(model_dir)
+    scores = model.score(prompt, candidate)    # output [score]
+    print(scores)