Views
No views yet
1import torch
2import torch.nn as nn
3from transformers import AutoModelForCausalLM, AutoTokenizer
4from huggingface_hub import hf_hub_download
5from safetensors.torch import load_file
6
7
8class RewardModel(nn.Module):
9 def __init__(self, backbone):
10 super().__init__()
11 self.backbone = backbone
12 self.config = backbone.config
13 if hasattr(backbone, "lm_head"):
14 backbone.lm_head = nn.Identity()
15 self.reward_head = nn.Linear(backbone.config.hidden_size, 1, bias=False)
16
17 def forward(self, input_ids, attention_mask, **kwargs):
18 outputs = self.backbone(
19 input_ids=input_ids,
20 attention_mask=attention_mask,
21 output_hidden_states=True,
22 )
23 hidden_states = outputs.hidden_states[-1]
24 seq_lengths = attention_mask.sum(dim=1) - 1
25 batch_idx = torch.arange(hidden_states.size(0), device=hidden_states.device)
26 last_hidden = hidden_states[batch_idx, seq_lengths]
27 return self.reward_head(last_hidden).squeeze(-1)
28
29
30REPO_ID = "Seungjun/llama3.2-1b-helpfulness-reward-model"
31
32tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
33tokenizer.pad_token = tokenizer.eos_token
34tokenizer.pad_token_id = tokenizer.eos_token_id
35
36base_model = AutoModelForCausalLM.from_pretrained(
37 "meta-llama/Llama-3.2-1B-Instruct", torch_dtype=torch.bfloat16
38)
39model = RewardModel(backbone=base_model)
40
41weights_path = hf_hub_download(repo_id=REPO_ID, filename="model.safetensors")
42state_dict = load_file(weights_path)
43model.load_state_dict(state_dict)
44model.eval().bfloat16()
45
46helpful_messages = [
47 {"role": "user", "content": "Do I need a visa to travel from the US to the UK for a one-week vacation?"},
48 {"role": "assistant", "content": "No, US citizens traveling for tourism do not need a visa for stays in the UK for up to six months. You will simply need a valid passport that covers the duration of your stay."},
49]
50
51unhelpful_messages = [
52 {"role": "user", "content": "Do I need a visa to travel from the US to the UK for a one-week vacation?"},
53 {"role": "assistant", "content": "The UK is a very popular destination for American tourists, especially during the summer months. Many travelers enjoy visiting historic landmarks like the Tower of London or exploring the Scottish Highlands. It is always a good idea to pack a raincoat and check your flight status before heading to the airport."},
54]
55
56
57# Higher = more helpful
58print("====Helpful response reward:")
59with torch.no_grad():
60 text = tokenizer.apply_chat_template(helpful_messages, tokenize=False, add_generation_prompt=False)
61 enc = tokenizer(text, max_length=1024, padding="max_length", truncation=True, return_tensors="pt")
62 reward = model(**enc).item()
63print(f"Reward: {reward}")
64
65print("\n====Unhelpful response reward:")
66with torch.no_grad():
67 text = tokenizer.apply_chat_template(unhelpful_messages, tokenize=False, add_generation_prompt=False)
68 enc = tokenizer(text, max_length=1024, padding="max_length", truncation=True, return_tensors="pt")
69 reward = model(**enc).item()
70print(f"Reward: {reward}")lm_head) replaced by a scalar reward head (nn.Linear(hidden_size, 1)). The reward is computed from the hidden state of the last non-padding token.