Update README.md

README.md

@@ -32,34 +32,31 @@ We build the PRM model as a LoRA adapter, which saves the memory to use it.
 As this LoRA adapter is built on `pixas/MedSSS_Policy`, you need to first prepare the base model in your platform.
 
 ```python
+from itertools import chain
+import torch
+from transformers import AutoModelForTokenClassification, AutoTokenizer
+from peft import PeftModel
 
 def obtain_prm_value_for_single_pair(tokenizer, value_model, inputs, outputs):
     # `outputs` generated by the MedSSS-Policy
-    response = outputs
-    completions = [f"Step" + completion if not completion.startswith("Step") else completion for k, completion in enumerate(outputs.split("\n\nStep"))]
-    
     messages = [
         {"role": "user", "content": inputs},
         {"role": "assistant", "content": response}
     ]
-    input_text = tokenizer.apply_chat_template(messages, tokenize=False)
-
-    response_begin_index = input_text.index(response)
+    
+    prompt_text = tokenizer.apply_chat_template(messages[:-1], tokenize=False, add_generation_prompt=True)
+    completions = ["Step" + completion if not completion.startswith("Step") else completion for completion in response.split("\n\nStep")]
 
-    pre_response_input = input_text[:response_begin_index]
-    after_response_input = input_text[response_begin_index + len(response):]
     completion_ids = [
         tokenizer(completion + "\n\n", add_special_tokens=False)['input_ids'] for completion in completions
     ]
-    
     response_id = list(chain(*completion_ids))
-    pre_response_id = tokenizer(pre_response_input, add_special_tokens=False)['input_ids']
-    after_response_id = tokenizer(after_response_input, add_special_tokens=False)['input_ids']
+    pre_response_id = tokenizer(prompt_text, add_special_tokens=False)['input_ids']
 
-    
-    input_ids = pre_response_id + response_id + after_response_id
-    
-    value = value_model(input_ids=torch.tensor(input_ids).unsqueeze(0).to(value_model.device))  # [1, N]
+    input_ids = pre_response_id + response_id
+
+    outputs = value_model(input_ids=torch.tensor(input_ids).unsqueeze(0).to(value_model.device))  # [1, N]
+    value = torch.softmax(outputs[0], dim=-1)[..., 1]
     
     completion_index = []
     for i, completion in enumerate(completion_ids):
@@ -70,12 +67,12 @@ def obtain_prm_value_for_single_pair(tokenizer, value_model, inputs, outputs):
     
     step_value = value[0, completion_index].cpu().numpy().tolist()
     return step_value
-from transformers import AutoModelForTokenClassification, AutoTokenizer
-from peft import PeftModel
-base_model = AutoModelForTokenClassification.from_pretrained("meta-llama/Llama-3.1-8B-Instruct",torch_dtype="auto",device_map="auto")
+
+
+base_model = AutoModelForTokenClassification.from_pretrained("pixas/MedSSS_Policy",torch_dtype="auto",device_map="auto")
 model = PeftModel.from_pretrained(base_model, "pixas/MedSSS_PRM", torc_dtype="auto", device_map="auto")
 tokenizer = AutoTokenizer.from_pretrained("pixas/MedSSS_PRM")
-steps
+
 input_text = "How to stop a cough?"
 step_wise_generation = "Step 0: Let's break down this problem step by step.\n\nStep 1: First [omitted]"