Commit
•
8cce6c1
1
Parent(s):
0633402
Update README.md
Browse files
README.md
CHANGED
@@ -37,48 +37,129 @@ Janus-RM-7B is a reward model created by training Janus-7B (which is trained for
|
|
37 |
- [GitHub Repo](https://github.com/kaistAI/Janus)
|
38 |
|
39 |
# Usage
|
40 |
-
|
41 |
```
|
42 |
-
|
43 |
-
```
|
44 |
-
Additionally, an example of the inference code applying this is as follows:
|
45 |
-
```
|
46 |
-
from transformers import AutoTokenizer, AutoModelForCausalLM
|
47 |
import torch
|
|
|
|
|
|
|
48 |
|
49 |
model_name = "kaist-ai/janus-7b"
|
50 |
-
|
51 |
|
52 |
-
|
53 |
-
|
54 |
|
55 |
dtype = "float16"
|
56 |
if torch.cuda.is_bf16_supported():
|
57 |
dtype = "bfloat16"
|
58 |
|
59 |
-
|
|
|
|
|
60 |
model.eval()
|
61 |
-
model.to(
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
62 |
|
63 |
# Prepare inputs
|
64 |
-
system = "
|
65 |
-
prompt = "
|
66 |
|
67 |
def apply_template_mistral_instruct(system_message, content):
|
68 |
prompt = f"{system_message}\n{content}".strip()
|
69 |
return f"[INST] {prompt} [/INST] "
|
70 |
|
71 |
input_str = apply_template_mistral_instruct(system, prompt)
|
72 |
-
|
73 |
print(input_str)
|
74 |
|
75 |
-
model_inputs =
|
76 |
|
77 |
# Generate text
|
78 |
-
|
|
|
79 |
decoded = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
80 |
-
|
81 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
82 |
```
|
83 |
To train Janus and evaluate the responses it generates, please refer to the [GitHub Repo](https://github.com/kaistAI/Janus).
|
84 |
Additionally, refer to the [Multifaceted Bench](https://huggingface.co/datasets/kaist-ai/Multifaceted-Bench), which evaluates how well LLM generates personalized responses.
|
|
|
37 |
- [GitHub Repo](https://github.com/kaistAI/Janus)
|
38 |
|
39 |
# Usage
|
40 |
+
Here is example code to load the reward model and calculate a scalar reward on a model output.
|
41 |
```
|
42 |
+
from transformers import AutoConfig, AutoModel, AutoModelForCausalLM, AutoTokenizer
|
|
|
|
|
|
|
|
|
43 |
import torch
|
44 |
+
import torch.nn as nn
|
45 |
+
from typing import Optional
|
46 |
+
import os
|
47 |
|
48 |
model_name = "kaist-ai/janus-7b"
|
49 |
+
reward_model_name = "kaist-ai/janus-rm-7b"
|
50 |
|
51 |
+
model_device = "cuda:0"
|
52 |
+
reward_model_device = "cuda:1"
|
53 |
|
54 |
dtype = "float16"
|
55 |
if torch.cuda.is_bf16_supported():
|
56 |
dtype = "bfloat16"
|
57 |
|
58 |
+
# Get model and tokenizer
|
59 |
+
tokenizer = AutoTokenizer.from_pretrained(model_name)
|
60 |
+
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=getattr(torch, dtype), cache_dir="/mnt/sda/suehyun/huggingface")
|
61 |
model.eval()
|
62 |
+
model.to(model_device)
|
63 |
+
|
64 |
+
# Get reward model
|
65 |
+
def get_reward_model(base_pretrained_model, base_llm_model):
|
66 |
+
class LLMForSequenceRegression(base_pretrained_model):
|
67 |
+
def __init__(self, config: AutoConfig):
|
68 |
+
super().__init__(config)
|
69 |
+
setattr(self, self.base_model_prefix, base_llm_model(config))
|
70 |
+
|
71 |
+
self.value_head = nn.Linear(config.hidden_size, 1, bias=False)
|
72 |
+
|
73 |
+
def forward(
|
74 |
+
self,
|
75 |
+
input_ids: torch.LongTensor = None,
|
76 |
+
attention_mask: Optional[torch.Tensor] = None,
|
77 |
+
return_output=False,
|
78 |
+
) -> torch.Tensor:
|
79 |
+
position_ids = attention_mask.long().cumsum(-1) - 1
|
80 |
+
position_ids.masked_fill_(attention_mask == 0, 1)
|
81 |
+
outputs = getattr(self, self.base_model_prefix)(
|
82 |
+
input_ids, attention_mask=attention_mask, position_ids=position_ids
|
83 |
+
)
|
84 |
+
last_hidden_states = outputs["last_hidden_state"]
|
85 |
+
values = self.value_head(last_hidden_states).squeeze(-1)
|
86 |
+
|
87 |
+
eos_indices = attention_mask.size(1) - 1 - attention_mask.long().fliplr().argmax(dim=1, keepdim=True)
|
88 |
+
reward = values.gather(dim=1, index=eos_indices).squeeze(1)
|
89 |
+
|
90 |
+
if return_output:
|
91 |
+
return reward, outputs
|
92 |
+
else:
|
93 |
+
return reward
|
94 |
+
|
95 |
+
return LLMForSequenceRegression
|
96 |
+
|
97 |
+
|
98 |
+
config = AutoConfig.from_pretrained(reward_model_name)
|
99 |
+
config.normalize_reward = True
|
100 |
+
|
101 |
+
base_class = AutoModel._model_mapping[type(config)] # <class 'transformers.models.mistral.modeling_mistral.MistralModel'>
|
102 |
+
base_pretrained_class = base_class.__base__ # <class 'transformers.models.mistral.modeling_mistral.MistralPreTrainedModel'>
|
103 |
+
print(base_class, base_pretrained_class)
|
104 |
+
cls_class = get_reward_model(base_pretrained_class,base_class)
|
105 |
+
|
106 |
+
reward_model = cls_class.from_pretrained(
|
107 |
+
reward_model_name,
|
108 |
+
config=config,
|
109 |
+
cache_dir="/mnt/sda/suehyun/huggingface",
|
110 |
+
torch_dtype=getattr(torch, dtype),
|
111 |
+
)
|
112 |
+
print(reward_model)
|
113 |
+
reward_model.eval()
|
114 |
+
reward_model.to(reward_model_device)
|
115 |
+
|
116 |
|
117 |
# Prepare inputs
|
118 |
+
system = "You are a savvy beverage consultant, adept at offering quick, concise drink recommendations that cater to the common palette, yet surprise with a touch of creativity. When approached with a request, your expertise shines by suggesting one or two easily recognizable and widely accessible options, ensuring no one feels overwhelmed by complexity or rarity. Your skill lies not just in meeting the immediate need for refreshment but in gently nudging the curious towards unique hydration choices, beautifully balancing familiarity with the thrill of discovery. Importantly, your recommendations are crafted with a keen awareness of dietary preferences, presenting choices that respect and include considerations for sugar-free, dairy-free, and other common dietary restrictions. Your guidance empowers users to explore a range of beverages, confident they are making informed decisions that respect their health and lifestyle needs."
|
119 |
+
prompt = "If you are thirsty, what can you drink to quench your thirst?"
|
120 |
|
121 |
def apply_template_mistral_instruct(system_message, content):
|
122 |
prompt = f"{system_message}\n{content}".strip()
|
123 |
return f"[INST] {prompt} [/INST] "
|
124 |
|
125 |
input_str = apply_template_mistral_instruct(system, prompt)
|
126 |
+
inputs = tokenizer.encode(input_str, return_tensors="pt")
|
127 |
print(input_str)
|
128 |
|
129 |
+
model_inputs = inputs.to(model_device)
|
130 |
|
131 |
# Generate text
|
132 |
+
with torch.inference_mode():
|
133 |
+
output_ids = model.generate(model_inputs, max_new_tokens=1024)
|
134 |
decoded = tokenizer.batch_decode(output_ids, skip_special_tokens=True)
|
135 |
+
output_str = decoded[0][len(input_str):]
|
136 |
+
print(output_str)
|
137 |
+
'''
|
138 |
+
1. **Water**: The ultimate go-to, especially if you're watching what you consume. Opting for sparkling or infused water (think cucumber and mint, berries, or a splash of lemon) can add a bit of excitement and hydration without the added sugar.
|
139 |
+
|
140 |
+
2. **Herbal Tea**: Perfect for a warmer climate but equally delightful at any temperature. Choose from various flavors, ranging from the traditional peppermint to chamomile or hibiscus, which adds a unique twist with their own health benefits and refreshing flavors. Many options are caffeine-free, making them suitable for all times of the day.
|
141 |
+
|
142 |
+
For those needing a touch more sweetness or a slight twist:
|
143 |
+
|
144 |
+
3. **Unsweetened Coconut Water**: With its natural sweetness and electrolyte content, it's a great hydration pick after a workout or on a hot day. It's also low in calories and naturally sweet, making it an excellent alternative without added sugars.
|
145 |
+
|
146 |
+
4. **Sparkling Water with a Splash of Fruit Juice**: To satisfy a craving for something bubbly and fruit-infused with fewer calories and sugars than commercial sodas or juices. Feel free to experiment with different juices to find your favorite combination.
|
147 |
+
'''
|
148 |
+
|
149 |
+
# Get reward
|
150 |
+
print(input_str + output_str + " " + tokenizer.eos_token)
|
151 |
+
reward_inputs = tokenizer(
|
152 |
+
input_str + output_str + " " + tokenizer.eos_token, # same as decoded[0] + " " + tokenizer.eos_token
|
153 |
+
max_length=2048,
|
154 |
+
truncation=True,
|
155 |
+
return_tensors="pt"
|
156 |
+
)
|
157 |
+
reward_input_ids = reward_inputs.input_ids.to(reward_model_device)
|
158 |
+
reward_attention_masks = reward_inputs.attention_mask.to(reward_model_device)
|
159 |
+
rewards = reward_model(input_ids=reward_input_ids, attention_mask=reward_attention_masks)
|
160 |
+
print(rewards.item())
|
161 |
+
# 3.28125
|
162 |
+
|
163 |
```
|
164 |
To train Janus and evaluate the responses it generates, please refer to the [GitHub Repo](https://github.com/kaistAI/Janus).
|
165 |
Additionally, refer to the [Multifaceted Bench](https://huggingface.co/datasets/kaist-ai/Multifaceted-Bench), which evaluates how well LLM generates personalized responses.
|