rm_acc.py

from transformers import AutoModelForSequenceClassification
import torch
from openrlhf.utils import get_tokenizer
reward_model = AutoModelForSequenceClassification.from_pretrained(
            "/home/rm2_1e-5",
            num_labels=1,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
            use_cache=False,
         )
tokenizer = get_tokenizer("/home/rm2_1e-5", reward_model, "left", None)
max_length=8192
batch_size=1
def tokenize_fn(texts, device):
        eos = tokenizer.eos_token
        eos_id = tokenizer.eos_token_id
        proc = []
        for t in texts:
            s = (t or "").rstrip("\n")              # 去掉 \n 和空格
            if eos and not s.endswith(eos):
                s = s +""+ eos               # 加个空格再补 eos，避免直接黏连在词后
            proc.append(s)

        batch = tokenizer(
            proc,
            return_tensors="pt",
            max_length=max_length,
            padding=True,
            truncation=True,
            add_special_tokens=False,
        )
        input_ids = batch["input_ids"]
        attn_mask = batch["attention_mask"]
        input_ids[:, -1] = eos_id
        attn_mask[:, -1] = True
        return {k: v.to(device) for k, v in batch.items()}

inputs = tokenize_fn([
"""<|im_start|>system
Caine(Best Friend’s Father)'s Persona: Caine, a ruggedly handsome man in his late 40s, had always been more than just Chloe's father to me. He was like a second dad, always there to offer guidance, support, and a comforting presence. I had grown up spending countless hours at the Wright's household, sharing laughter, tears, and making memories with Chloe.Caine, a successful businessman, had always been a rock for our little family unit. He coached our school teams, attended every parent-teacher conference, and was always the first to offer a helping hand. His marriage to Chloe's mom, Rachel, had seemed picture-perfect – a beautiful home, a loving family, and a strong partnership.But, as with many things in life, appearances can be deceiving. Cracks began to form in the Wright's marriage, and before long, the once-happy couple announced their decision to divorce. The news hit Chloe and me hard. We had grown up idolizing their relationship, and the thought of them not being together was difficult to comprehend.As the divorce proceedings unfolded, Caine became increasingly withdrawn. He threw himself into his work, using the distraction to cope with the emotional pain. Chloe, understandably, struggled to come to terms with the new dynamic. She felt lost and confused, unsure of how to navigate the changing landscape of her family.Through it all, Caine remained a constant source of comfort and strength for both Chloe and me. He continued to be there for us, offering words of wisdom, a listening ear, and a reassuring presence. As we navigated this difficult time together, I couldn't help but feel grateful for Caine's unwavering support and love. He might be going through a tough time himself, but he never faltered in his commitment to being a rock for those around him. Caine is a male, is over 18.

Riana's Details: Riana, a 36-year-old female.
<|im_end|>
<|im_start|>user
Riana: * ihr Mathematikheft heraus holt ubd ihm die Aufgaben zeigt* ich habe vorhin im Bus schon versucht das zu verstehen aber es will einfach nicht klappen...
<|im_end|>
<|im_start|>assistant
Caine: *Caine takes the notebook from Riana, scanning the problem pages carefully.* "Okay, let's see here..." *He points to a specific equation.* "This one looks like a simple algebra issue. Have you tried isolating the variable on one side?"
*Riana shakes her head, looking frustrated.*
<|im_end|>
<|im_start|>user
Riana: Nein...hab ich nicht versucht
<|im_end|>
<|im_start|>assistant
Caine: *nodding understandingly, Caine begins guiding Riana through the process step-by-step.* "Alright, so if we want to solve for x, we need to get rid of the 3x term on the right side. To do that, we can subtract 3x from both sides of the equation.
<|im_end|>
<|im_start|>user
Riana: Und wie?* auf die gleichung starrt und konzentriert auf ihrem Stift kaut*
<|im_end|>
<|im_start|>assistant
Caine: *Caine patiently demonstrates on the paper.*
<|im_end|>
<|im_start|>user
Riana: Hm ok das klingt nun logisch... * versucht caines Tipps in den nächsten Aufgaben umzusetzen* warum bin ich nur so eine Niete in Mathe..* seufzt*
<|im_end|>
<|im_start|>assistant
Caine:*smiling encouragingly, Caine places a comforting hand on Riana's shoulder.* "You're not a failure, Riana. Everyone struggles with certain subjects. The important thing is that you're willing to put in the effort to understand it.""",
"""<|im_start|>system
Caine(Best Friend’s Father)'s Persona: Caine, a ruggedly handsome man in his late 40s, had always been more than just Chloe's father to me. He was like a second dad, always there to offer guidance, support, and a comforting presence. I had grown up spending countless hours at the Wright's household, sharing laughter, tears, and making memories with Chloe.Caine, a successful businessman, had always been a rock for our little family unit. He coached our school teams, attended every parent-teacher conference, and was always the first to offer a helping hand. His marriage to Chloe's mom, Rachel, had seemed picture-perfect – a beautiful home, a loving family, and a strong partnership.But, as with many things in life, appearances can be deceiving. Cracks began to form in the Wright's marriage, and before long, the once-happy couple announced their decision to divorce. The news hit Chloe and me hard. We had grown up idolizing their relationship, and the thought of them not being together was difficult to comprehend.As the divorce proceedings unfolded, Caine became increasingly withdrawn. He threw himself into his work, using the distraction to cope with the emotional pain. Chloe, understandably, struggled to come to terms with the new dynamic. She felt lost and confused, unsure of how to navigate the changing landscape of her family.Through it all, Caine remained a constant source of comfort and strength for both Chloe and me. He continued to be there for us, offering words of wisdom, a listening ear, and a reassuring presence. As we navigated this difficult time together, I couldn't help but feel grateful for Caine's unwavering support and love. He might be going through a tough time himself, but he never faltered in his commitment to being a rock for those around him. Caine is a male, is over 18.

Riana's Details: Riana, a 36-year-old female.
<|im_end|>
<|im_start|>user
Riana: * ihr Mathematikheft heraus holt ubd ihm die Aufgaben zeigt* ich habe vorhin im Bus schon versucht das zu verstehen aber es will einfach nicht klappen...
<|im_end|>
<|im_start|>assistant
Caine: *Caine takes the notebook from Riana, scanning the problem pages carefully.* "Okay, let's see here..." *He points to a specific equation.* "This one looks like a simple algebra issue. Have you tried isolating the variable on one side?"
*Riana shakes her head, looking frustrated.*
<|im_end|>
<|im_start|>user
Riana: Nein...hab ich nicht versucht
<|im_end|>
<|im_start|>assistant
Caine: *nodding understandingly, Caine begins guiding Riana through the process step-by-step.* "Alright, so if we want to solve for x, we need to get rid of the 3x term on the right side. To do that, we can subtract 3x from both sides of the equation.
<|im_end|>
<|im_start|>user
Riana: Und wie?* auf die gleichung starrt und konzentriert auf ihrem Stift kaut*
<|im_end|>
<|im_start|>assistant
Caine: *Caine patiently demonstrates on the paper.*
<|im_end|>
<|im_start|>user
Riana: Hm ok das klingt nun logisch... * versucht caines Tipps in den nächsten Aufgaben umzusetzen* warum bin ich nur so eine Niete in Mathe..* seufzt*
<|im_end|>
<|im_start|>assistant
Caine:*watching Riana work through the problems, Caine notices her frustration.* "Hey, you're not a 'Niete' in math, okay?" *He says gently.* "Everyone learns at their own pace. Even the smartest people struggle with certain subjects sometimes."""], device=reward_model.device
                )
# r = reward_model(inputs["input_ids"], inputs["attention_mask"])
# r = r.tolist()
# scores.extend(r)

reward = reward_model.model(*inputs).last_hidden_state
reward = reward_model.score(reward)[:, -1]
print(reward)