File size: 3,146 Bytes
fa36636 08d8702 fa36636 133e0e7 fa36636 08d8702 fa36636 08d8702 fa36636 4bd4fd1 fa36636 4bd4fd1 fa36636 4bd4fd1 fa36636 4bd4fd1 fa36636 4bd4fd1 fa36636 4bd4fd1 fa36636 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 |
---
license: mit
language:
- en
base_model:
- pixas/MedSSS_Policy
pipeline_tag: token-classification
---
<div align="center">
<h1>
MedSSS-8B-PRM
</h1>
</div>
<div align="center">
<a href="https://github.com/pixas/MedSSS" target="_blank">GitHub</a> | <a href="https://arxiv.org/abs/2501.12051" target="_blank">Paper</a>
</div>
# <span>Introduction</span>
**MedSSS-PRM** is trained with the newly proposed soft dual-sided object, designed for identifying intermediate erroneous steps within a correct medical reasoning trajectory.
It will assign a `[0-1]` float value for every internal reasoning step of **MedSSS-Policy**.
For more information, visit our GitHub repository:
[https://github.com/pixas/MedSSS](https://github.com/pixas/MedSSS).
# <span>Usage</span>
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
messages = [
{"role": "user", "content": inputs},
{"role": "assistant", "content": 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")]
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(prompt_text, add_special_tokens=False)['input_ids']
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):
if i == 0:
completion_index.append(len(completion) + len(pre_response_id) - 1)
else:
completion_index.append(completion_index[-1] + len(completion))
step_value = value[0, completion_index].cpu().numpy().tolist()
return step_value
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")
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]"
value = obtain_prm_value_for_single_pair(tokenizer, model, input_text, step_wise_generation)
print(value)
```
MedSSS-PRM uses "\n\nStep" to separate intermediate steps. So the token classification happens before the next "Step k: " or the end of the sequence. |