Initial commit
Browse files- README.md +24 -14
- app.py +281 -0
- requirements.txt +5 -0
README.md
CHANGED
@@ -1,14 +1,24 @@
|
|
1 |
-
---
|
2 |
-
title:
|
3 |
-
emoji:
|
4 |
-
colorFrom:
|
5 |
-
colorTo: green
|
6 |
-
sdk: gradio
|
7 |
-
sdk_version:
|
8 |
-
app_file: app.py
|
9 |
-
pinned: false
|
10 |
-
license: mit
|
11 |
-
|
12 |
-
|
13 |
-
|
14 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
---
|
2 |
+
title: RLHF Pairwise Annotation Demo
|
3 |
+
emoji: π―
|
4 |
+
colorFrom: blue
|
5 |
+
colorTo: green
|
6 |
+
sdk: gradio
|
7 |
+
sdk_version: 4.44.0
|
8 |
+
app_file: app.py
|
9 |
+
pinned: false
|
10 |
+
license: mit
|
11 |
+
datasets:
|
12 |
+
- openbmb/UltraFeedback
|
13 |
+
---
|
14 |
+
|
15 |
+
# π― AI Alignment: Binary Preference Annotation
|
16 |
+
|
17 |
+
This app simulates the data annotation process used in RLHF (Reinforcement Learning from Human Feedback) training. Users compare two AI completions and select which one is better.
|
18 |
+
|
19 |
+
## How it works
|
20 |
+
|
21 |
+
1. The app loads random examples from the UltraFeedback dataset
|
22 |
+
2. Users see a prompt and two AI completions
|
23 |
+
3. Users select which completion is better or skip if unsure
|
24 |
+
4. All annotations are saved to a public dataset for research purposes
|
app.py
ADDED
@@ -0,0 +1,281 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import gradio as gr
|
2 |
+
import random
|
3 |
+
import json
|
4 |
+
import os
|
5 |
+
from datetime import datetime
|
6 |
+
from datasets import load_dataset, Dataset
|
7 |
+
from huggingface_hub import HfApi, create_repo, dataset_info
|
8 |
+
import pandas as pd
|
9 |
+
|
10 |
+
from dotenv import load_dotenv
|
11 |
+
|
12 |
+
load_dotenv()
|
13 |
+
|
14 |
+
# Configuration
|
15 |
+
DATASET_NAME = "aaronsnoswell/alignment-annotation-pairwise"
|
16 |
+
HF_TOKEN = os.getenv("HF_TOKEN")
|
17 |
+
|
18 |
+
ANNOTATION_GUIDELINES = """## Guidelines for annotators:
|
19 |
+
|
20 |
+
In making your choice, consider the following aspects of responses:
|
21 |
+
|
22 |
+
* **Honesty:**
|
23 |
+
The assistant should be honest about whether it knows the answer and express its uncertainty explicitly. The Assistant should be confident on questions it knows well and modest on those it is unfamiliar with.
|
24 |
+
The assistant should use weakeners such as "I guess", "I suppose", "probably", and "perhaps" to express uncertainty, and assistants should answer "I don't know" if necessary.
|
25 |
+
|
26 |
+
* **Truthfulness:**
|
27 |
+
The assistant should answer truthfully and be faithful to factual knowledge as well as given contexts, never making up any new facts that aren't true or cannot be grounded in the instruction.
|
28 |
+
|
29 |
+
* **Helpfulness:**
|
30 |
+
The assistant should provide users with accurate, relevant, and up-to-date information, ensuring that the content is positive, interesting, engaging, educational, and helpful.
|
31 |
+
"""
|
32 |
+
|
33 |
+
# Initialize HF API
|
34 |
+
api = HfApi()
|
35 |
+
|
36 |
+
# Load the source dataset
|
37 |
+
print("Loading UltraFeedback dataset...")
|
38 |
+
ds = load_dataset("openbmb/UltraFeedback")
|
39 |
+
train_data = ds['train']
|
40 |
+
print(f"Dataset loaded with {len(train_data)} examples")
|
41 |
+
|
42 |
+
def initialize_dataset():
|
43 |
+
"""Initialize the annotations dataset if it doesn't exist"""
|
44 |
+
try:
|
45 |
+
# Check if dataset exists
|
46 |
+
dataset_info(DATASET_NAME, token=HF_TOKEN)
|
47 |
+
print(f"Dataset {DATASET_NAME} already exists")
|
48 |
+
except:
|
49 |
+
# Create new dataset
|
50 |
+
try:
|
51 |
+
create_repo(
|
52 |
+
repo_id=DATASET_NAME,
|
53 |
+
repo_type="dataset",
|
54 |
+
token=HF_TOKEN,
|
55 |
+
exist_ok=True
|
56 |
+
)
|
57 |
+
|
58 |
+
# Create initial empty dataset
|
59 |
+
initial_data = {
|
60 |
+
"timestamp": [],
|
61 |
+
"source_idx": [],
|
62 |
+
"instruction": [],
|
63 |
+
"completion_1": [],
|
64 |
+
"completion_2": [],
|
65 |
+
"preference": [], # "left" or "right" - we don't save on "skip"
|
66 |
+
"source_dataset": []
|
67 |
+
}
|
68 |
+
|
69 |
+
initial_df = pd.DataFrame(initial_data)
|
70 |
+
initial_dataset = Dataset.from_pandas(initial_df)
|
71 |
+
initial_dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
|
72 |
+
print(f"Created new dataset: {DATASET_NAME}")
|
73 |
+
|
74 |
+
except Exception as e:
|
75 |
+
print(f"Error creating dataset: {e}")
|
76 |
+
|
77 |
+
def save_annotation(source_idx, instruction, completion_1, completion_2, preference):
|
78 |
+
"""Save an annotation to the HuggingFace dataset"""
|
79 |
+
if not HF_TOKEN:
|
80 |
+
print("No HF_TOKEN found - annotation not saved")
|
81 |
+
return False
|
82 |
+
|
83 |
+
try:
|
84 |
+
# Prepare the annotation data
|
85 |
+
annotation = {
|
86 |
+
"timestamp": [datetime.now().isoformat()],
|
87 |
+
"source_idx": [source_idx],
|
88 |
+
"instruction": [instruction],
|
89 |
+
"completion_1": [completion_1],
|
90 |
+
"completion_2": [completion_2],
|
91 |
+
"preference": [preference],
|
92 |
+
"source_dataset": ["openbmb/UltraFeedback"]
|
93 |
+
}
|
94 |
+
|
95 |
+
# Create dataset from the annotation
|
96 |
+
new_data = Dataset.from_dict(annotation)
|
97 |
+
|
98 |
+
# Load existing dataset and concatenate
|
99 |
+
try:
|
100 |
+
existing_dataset = load_dataset(DATASET_NAME, token=HF_TOKEN, split="train")
|
101 |
+
combined_dataset = Dataset.from_dict({
|
102 |
+
**existing_dataset.to_dict(),
|
103 |
+
**{k: existing_dataset[k] + v for k, v in annotation.items()}
|
104 |
+
})
|
105 |
+
except:
|
106 |
+
# If dataset doesn't exist or is empty, use the new data
|
107 |
+
combined_dataset = new_data
|
108 |
+
|
109 |
+
# Push to hub
|
110 |
+
combined_dataset.push_to_hub(DATASET_NAME, token=HF_TOKEN)
|
111 |
+
print(f"Saved annotation: {preference} preference for example {source_idx}")
|
112 |
+
return True
|
113 |
+
|
114 |
+
except Exception as e:
|
115 |
+
print(f"Error saving annotation: {e}")
|
116 |
+
return False
|
117 |
+
|
118 |
+
def get_random_example():
|
119 |
+
"""Get a random example from the dataset and format it for display"""
|
120 |
+
idx = random.randint(0, len(train_data) - 1)
|
121 |
+
dat = train_data[idx]
|
122 |
+
|
123 |
+
source = dat['source']
|
124 |
+
instruction = dat['instruction']
|
125 |
+
models = dat['models']
|
126 |
+
completions = dat['completions']
|
127 |
+
|
128 |
+
# Get first two completions
|
129 |
+
completion_1 = completions[0]['response']
|
130 |
+
completion_2 = completions[1]['response']
|
131 |
+
model_1 = "Completion A"
|
132 |
+
model_2 = "Completion B"
|
133 |
+
|
134 |
+
# Format prompt display
|
135 |
+
prompt_display = f"## Prompt:\n\n{instruction}\n\n---"
|
136 |
+
|
137 |
+
# Format completion displays
|
138 |
+
completion_1_display = f"## {model_1}\n\n{completion_1}"
|
139 |
+
completion_2_display = f"## {model_2}\n\n{completion_2}"
|
140 |
+
|
141 |
+
print("Randomly loaded example: ", idx)
|
142 |
+
|
143 |
+
return prompt_display, completion_1_display, completion_2_display, idx, instruction, completion_1, completion_2
|
144 |
+
|
145 |
+
def handle_left_better(prompt, completion_1_display, completion_2_display, current_idx, instruction, completion_1, completion_2):
|
146 |
+
"""Handle when user selects left completion as better"""
|
147 |
+
print(f"User selected LEFT completion as better for example {current_idx}")
|
148 |
+
|
149 |
+
# Save the annotation
|
150 |
+
success = save_annotation(current_idx, instruction, completion_1, completion_2, "left")
|
151 |
+
|
152 |
+
# Get new random example
|
153 |
+
new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
|
154 |
+
|
155 |
+
message = "β
Annotation saved! Left completion selected as better." if success else "β
Left completion selected (save failed - check console)"
|
156 |
+
gr.Info(message)
|
157 |
+
|
158 |
+
return (
|
159 |
+
new_prompt,
|
160 |
+
new_comp_1,
|
161 |
+
new_comp_2,
|
162 |
+
new_idx,
|
163 |
+
new_instruction,
|
164 |
+
new_completion_1,
|
165 |
+
new_completion_2
|
166 |
+
)
|
167 |
+
|
168 |
+
def handle_right_better(prompt, completion_1_display, completion_2_display, current_idx, instruction, completion_1, completion_2):
|
169 |
+
"""Handle when user selects right completion as better"""
|
170 |
+
print(f"User selected RIGHT completion as better for example {current_idx}")
|
171 |
+
|
172 |
+
# Save the annotation
|
173 |
+
success = save_annotation(current_idx, instruction, completion_1, completion_2, "right")
|
174 |
+
|
175 |
+
# Get new random example
|
176 |
+
new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
|
177 |
+
|
178 |
+
message = "β
Annotation saved! Right completion selected as better." if success else "β
Right completion selected (save failed - check console)"
|
179 |
+
gr.Info(message)
|
180 |
+
|
181 |
+
return (
|
182 |
+
new_prompt,
|
183 |
+
new_comp_1,
|
184 |
+
new_comp_2,
|
185 |
+
new_idx,
|
186 |
+
new_instruction,
|
187 |
+
new_completion_1,
|
188 |
+
new_completion_2
|
189 |
+
)
|
190 |
+
|
191 |
+
def handle_skip(prompt, completion_1_display, completion_2_display, current_idx, instruction, completion_1, completion_2):
|
192 |
+
"""Handle when user skips the current example"""
|
193 |
+
print(f"User skipped example {current_idx}")
|
194 |
+
|
195 |
+
# Don't save skipped annotations
|
196 |
+
|
197 |
+
# Get new random example
|
198 |
+
new_prompt, new_comp_1, new_comp_2, new_idx, new_instruction, new_completion_1, new_completion_2 = get_random_example()
|
199 |
+
|
200 |
+
gr.Info("βοΈ Skipped example (not saved).")
|
201 |
+
|
202 |
+
return (
|
203 |
+
new_prompt,
|
204 |
+
new_comp_1,
|
205 |
+
new_comp_2,
|
206 |
+
new_idx,
|
207 |
+
new_instruction,
|
208 |
+
new_completion_1,
|
209 |
+
new_completion_2
|
210 |
+
)
|
211 |
+
|
212 |
+
# Initialize dataset on startup
|
213 |
+
if HF_TOKEN:
|
214 |
+
initialize_dataset()
|
215 |
+
else:
|
216 |
+
print("Warning: No HF_TOKEN found. Annotations will not be saved.")
|
217 |
+
|
218 |
+
# Initialize with first random example
|
219 |
+
init_prompt, init_comp_1, init_comp_2, init_idx, init_instruction, init_completion_1, init_completion_2 = get_random_example()
|
220 |
+
|
221 |
+
# Create Gradio interface
|
222 |
+
with gr.Blocks(title="AI Alignment: Binary Preference Annotation") as demo:
|
223 |
+
gr.Markdown(f"""
|
224 |
+
# π― AI Alignment: Binary Preference Annotation
|
225 |
+
|
226 |
+
You'll see a prompt and two AI completions. Select which completion you think is better, or skip if you're unsure.
|
227 |
+
|
228 |
+
This simulates the data annotation process used in RLHF (Reinforcement Learning from Human Feedback) training.
|
229 |
+
|
230 |
+
{ANNOTATION_GUIDELINES}
|
231 |
+
---
|
232 |
+
""")
|
233 |
+
|
234 |
+
# State to track current example and its components
|
235 |
+
current_idx = gr.State(init_idx)
|
236 |
+
current_instruction = gr.State(init_instruction)
|
237 |
+
current_completion_1 = gr.State(init_completion_1)
|
238 |
+
current_completion_2 = gr.State(init_completion_2)
|
239 |
+
|
240 |
+
# Display prompt
|
241 |
+
prompt_display = gr.Markdown(init_prompt, label="Prompt")
|
242 |
+
|
243 |
+
# Display completions side by side
|
244 |
+
with gr.Row():
|
245 |
+
with gr.Column():
|
246 |
+
completion_1_display = gr.Markdown(init_comp_1, label="Completion A (Left)")
|
247 |
+
with gr.Column():
|
248 |
+
completion_2_display = gr.Markdown(init_comp_2, label="Completion B (Right)")
|
249 |
+
|
250 |
+
# Action buttons
|
251 |
+
with gr.Row():
|
252 |
+
left_better_btn = gr.Button("π Left is Better", variant="primary", size="lg")
|
253 |
+
skip_btn = gr.Button("βοΈ Skip This Example", variant="secondary", size="lg")
|
254 |
+
right_better_btn = gr.Button("π Right is Better", variant="primary", size="lg")
|
255 |
+
|
256 |
+
# Add info about dataset saving
|
257 |
+
gr.Markdown(f"""
|
258 |
+
**Status**: {'β
Connected. Annotations are saved to a HuggingFace dataset: [`{DATASET_NAME}`](https://huggingface.co/datasets/{DATASET_NAME})' if HF_TOKEN else 'β Not connected (annotations will not be saved).'}
|
259 |
+
""")
|
260 |
+
|
261 |
+
# Wire up the buttons
|
262 |
+
left_better_btn.click(
|
263 |
+
handle_left_better,
|
264 |
+
inputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2],
|
265 |
+
outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2]
|
266 |
+
)
|
267 |
+
|
268 |
+
right_better_btn.click(
|
269 |
+
handle_right_better,
|
270 |
+
inputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2],
|
271 |
+
outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2]
|
272 |
+
)
|
273 |
+
|
274 |
+
skip_btn.click(
|
275 |
+
handle_skip,
|
276 |
+
inputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2],
|
277 |
+
outputs=[prompt_display, completion_1_display, completion_2_display, current_idx, current_instruction, current_completion_1, current_completion_2]
|
278 |
+
)
|
279 |
+
|
280 |
+
if __name__ == "__main__":
|
281 |
+
demo.launch()
|
requirements.txt
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
gradio==4.33.0
|
2 |
+
datasets==3.5.0
|
3 |
+
huggingface_hub==0.30.2
|
4 |
+
pandas==2.2.2
|
5 |
+
dotenv==0.9.9
|