aaronsnoswell commited on
Commit
7b4b619
Β·
verified Β·
1 Parent(s): d7403d6

Initial commit

Browse files
Files changed (3) hide show
  1. README.md +24 -14
  2. app.py +281 -0
  3. requirements.txt +5 -0
README.md CHANGED
@@ -1,14 +1,24 @@
1
- ---
2
- title: Alignment Annotation Pairwise
3
- emoji: πŸƒ
4
- colorFrom: indigo
5
- colorTo: green
6
- sdk: gradio
7
- sdk_version: 5.39.0
8
- app_file: app.py
9
- pinned: false
10
- license: mit
11
- short_description: Try out pairwise data annotation for alignment
12
- ---
13
-
14
- Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 
 
 
 
 
 
 
 
 
 
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