Setting the file. One moment.
Dataset Inspector · Huggingface LLM Trainer · huggingface/skills · Skills Docs
Back to the top of the page 14.10
Unsloth
202
def main
— line 202
This file
Number 14.12
Position 12 of 18
Type Python
Size 15 KB
Lines 417 scripts/ dataset_inspector.py
Python · 417 lines · 15 KB
17 "script_args": ["--dataset", "your/dataset", "--split", "train"]
18 })
19 """
20
21 import argparse
22 import sys
23 import json
24 import urllib.request
25 import urllib.parse
26 from typing import List, Dict, Any
27
28
29 def parse_args ():
30 parser = argparse.ArgumentParser( description = "Inspect dataset format for TRL training" )
31 parser.add_argument( "--dataset" , type = str , required = True , help = "Dataset name" )
32 parser.add_argument( "--split" , type = str , default = "train" , help = "Dataset split (default: train)" )
33 parser.add_argument( "--config" , type = str , default = "default" , help = "Dataset config name (default: default)" )
34 parser.add_argument( "--preview" , type = int , default = 150 , help = "Max chars per field preview" )
35 parser.add_argument( "--samples" , type = int , default = 5 , help = "Number of samples to fetch (default: 5)" )
36 parser.add_argument( "--json-output" , action = "store_true" , help = "Output as JSON" )
37 return parser.parse_args()
38
39
40 def api_request (url: str ) -> Dict:
41 """Make API request to Datasets Server"""
42 try :
43 with urllib.request.urlopen(url, timeout = 10 ) as response:
44 return json.loads(response.read().decode())
45 except urllib.error.HTTPError as e:
46 if e.code == 404 :
47 return None
48 raise Exception ( f "API request failed: { e.code } { e.reason } " )
49 except Exception as e:
50 raise Exception ( f "API request failed: { str (e) } " )
51
52
53 def get_splits (dataset: str ) -> Dict:
54 """Get available splits for dataset"""
55 url = f "https://datasets-server.huggingface.co/splits?dataset= { urllib.parse.quote(dataset) } "
56 return api_request(url)
57
58
59 def get_rows (dataset: str , config: str , split: str , offset: int = 0 , length: int = 5 ) -> Dict:
60 """Get rows from dataset"""
61 url = f "https://datasets-server.huggingface.co/rows?dataset= { urllib.parse.quote(dataset) } &config= { config } &split= { split } &offset= { offset } &length= { length } "
62 return api_request(url)
63
64
65 def find_columns (columns: List[ str ], patterns: List[ str ]) -> List[ str ]:
66 """Find columns matching patterns"""
67 return [c for c in columns if any (p in c.lower() for p in patterns)]
68
69
70 def check_sft_compatibility (columns: List[ str ]) -> Dict[ str , Any]:
71 """Check SFT compatibility"""
72 has_messages = "messages" in columns
73 has_text = "text" in columns
74 has_prompt_completion = "prompt" in columns and "completion" in columns
75
76 ready = has_messages or has_text or has_prompt_completion
77
78 possible_prompt = find_columns(columns, [ "prompt" , "instruction" , "question" , "input" ])
79 possible_response = find_columns(columns, [ "response" , "completion" , "output" , "answer" ])
80
81 return {
82 "ready" : ready,
83 "reason" : "messages" if has_messages else "text" if has_text else "prompt+completion" if has_prompt_completion else None ,
84 "possible_prompt" : possible_prompt[ 0 ] if possible_prompt else None ,
85 "possible_response" : possible_response[ 0 ] if possible_response else None ,
86 "has_context" : "context" in columns,
87 }
88
89
90 def check_dpo_compatibility (columns: List[ str ]) -> Dict[ str , Any]:
91 """Check DPO compatibility"""
92 has_standard = "prompt" in columns and "chosen" in columns and "rejected" in columns
93
94 possible_prompt = find_columns(columns, [ "prompt" , "instruction" , "question" , "input" ])
95 possible_chosen = find_columns(columns, [ "chosen" , "preferred" , "winner" ])
96 possible_rejected = find_columns(columns, [ "rejected" , "dispreferred" , "loser" ])
97
98 can_map = bool (possible_prompt and possible_chosen and possible_rejected)
99
100 return {
101 "ready" : has_standard,
102 "can_map" : can_map,
103 "prompt_col" : possible_prompt[ 0 ] if possible_prompt else None ,
104 "chosen_col" : possible_chosen[ 0 ] if possible_chosen else None ,
105 "rejected_col" : possible_rejected[ 0 ] if possible_rejected else None ,
106 }
107
108
109 def check_grpo_compatibility (columns: List[ str ]) -> Dict[ str , Any]:
110 """Check GRPO compatibility"""
111 has_prompt = "prompt" in columns
112 has_no_responses = "chosen" not in columns and "rejected" not in columns
113
114 possible_prompt = find_columns(columns, [ "prompt" , "instruction" , "question" , "input" ])
115
116 return {
117 "ready" : has_prompt and has_no_responses,
118 "can_map" : bool (possible_prompt) and has_no_responses,
119 "prompt_col" : possible_prompt[ 0 ] if possible_prompt else None ,
120 }
121
122
123 def check_kto_compatibility (columns: List[ str ]) -> Dict[ str , Any]:
124 """Check KTO compatibility"""
125 return { "ready" : "prompt" in columns and "completion" in columns and "label" in columns}
126
127
128 def generate_mapping_code (method: str , info: Dict[ str , Any]) -> str :
129 """Generate mapping code for a training method"""
130 if method == "SFT" :
131 if info[ "ready" ]:
132 return None
133
134 prompt_col = info.get( "possible_prompt" )
135 response_col = info.get( "possible_response" )
136 has_context = info.get( "has_context" , False )
137
138 if not prompt_col:
139 return None
140
141 if has_context and response_col:
142 return f """def format_for_sft(example):
143 text = f"Instruction: {{ example[' { prompt_col } '] }}\n\n "
144 if example.get('context'):
145 text += f"Context: {{ example['context'] }}\n\n "
146 text += f"Response: {{ example[' { response_col } '] }} "
147 return {{ 'text': text }}
148
149 dataset = dataset.map(format_for_sft, remove_columns=dataset.column_names)"""
150 elif response_col:
151 return f """def format_for_sft(example):
152 return {{ 'text': f" {{ example[' { prompt_col } '] }}\n\n{{ example[' { response_col } '] }}}}
153
154 dataset = dataset.map(format_for_sft, remove_columns=dataset.column_names)"""
155 else :
156 return f """def format_for_sft(example):
157 return {{ 'text': example[' { prompt_col } '] }}
158
159 dataset = dataset.map(format_for_sft, remove_columns=dataset.column_names)"""
160
161 elif method == "DPO" :
162 if info[ "ready" ] or not info[ "can_map" ]:
163 return None
164
165 return f """def format_for_dpo(example):
166 return {{
167 'prompt': example[' { info[ 'prompt_col' ] } '],
168 'chosen': example[' { info[ 'chosen_col' ] } '],
169 'rejected': example[' { info[ 'rejected_col' ] } '],
170 }}
171
172 dataset = dataset.map(format_for_dpo, remove_columns=dataset.column_names)"""
173
174 elif method == "GRPO" :
175 if info[ "ready" ] or not info[ "can_map" ]:
176 return None
177
178 return f """def format_for_grpo(example):
179 return {{ 'prompt': example[' { info[ 'prompt_col' ] } '] }}
180
181 dataset = dataset.map(format_for_grpo, remove_columns=dataset.column_names)"""
182
183 return None
184
185
186 def format_value_preview (value: Any, max_chars: int ) -> str :
187 """Format value for preview"""
188 if value is None :
189 return "None"
190 elif isinstance (value, str ):
191 return value[:max_chars] + ( "..." if len (value) > max_chars else "" )
192 elif isinstance (value, list ):
193 if len (value) > 0 and isinstance (value[ 0 ], dict ):
194 return f "[ { len (value) } items] Keys: { list (value[ 0 ].keys()) } "
195 preview = str (value)
196 return preview[:max_chars] + ( "..." if len (preview) > max_chars else "" )
197 else :
198 preview = str (value)
199 return preview[:max_chars] + ( "..." if len (preview) > max_chars else "" )
200
201
202 def main ():
203 args = parse_args()
204
205 print ( f "Fetching dataset info via Datasets Server API..." )
206
207 try :
208 # Get splits info
209 splits_data = get_splits(args.dataset)
210 if not splits_data or "splits" not in splits_data:
211 print ( f "ERROR: Could not fetch splits for dataset ' { args.dataset } '" )
212 print ( f " Dataset may not exist or is not accessible via Datasets Server API" )
213 sys.exit( 1 )
214
215 # Find the right config
216 available_configs = set ()
217 split_found = False
218 config_to_use = args.config
219
220 for split_info in splits_data[ "splits" ]:
221 available_configs.add(split_info[ "config" ])
222 if split_info[ "config" ] == args.config and split_info[ "split" ] == args.split:
223 split_found = True
224
225 # If default config not found, try first available
226 if not split_found and available_configs:
227 config_to_use = list (available_configs)[ 0 ]
228 print ( f "Config ' { args.config } ' not found, trying ' { config_to_use } '..." )
229
230 # Get rows
231 rows_data = get_rows(args.dataset, config_to_use, args.split, offset = 0 , length = args.samples)
232
233 if not rows_data or "rows" not in rows_data:
234 print ( f "ERROR: Could not fetch rows for dataset ' { args.dataset } '" )
235 print ( f " Split ' { args.split } ' may not exist" )
236 print ( f " Available configs: { ', ' .join( sorted (available_configs)) } " )
237 sys.exit( 1 )
238
239 rows = rows_data[ "rows" ]
240 if not rows:
241 print ( f "ERROR: No rows found in split ' { args.split } '" )
242 sys.exit( 1 )
243
244 # Extract column info from first row
245 first_row = rows[ 0 ][ "row" ]
246 columns = list (first_row.keys())
247 features = rows_data.get( "features" , [])
248
249 # Get total count if available
250 total_examples = "Unknown"
251 for split_info in splits_data[ "splits" ]:
252 if split_info[ "config" ] == config_to_use and split_info[ "split" ] == args.split:
253 total_examples = f " { split_info.get( 'num_examples' , 'Unknown' ) :,} " if isinstance (split_info.get( 'num_examples' ), int ) else "Unknown"
254 break
255
256 except Exception as e:
257 print ( f "ERROR: { str (e) } " )
258 sys.exit( 1 )
259
260 # Run compatibility checks
261 sft_info = check_sft_compatibility(columns)
262 dpo_info = check_dpo_compatibility(columns)
263 grpo_info = check_grpo_compatibility(columns)
264 kto_info = check_kto_compatibility(columns)
265
266 # Determine recommended methods
267 recommended = []
268 if sft_info[ "ready" ]:
269 recommended.append( "SFT" )
270 elif sft_info[ "possible_prompt" ]:
271 recommended.append( "SFT (needs mapping)" )
272
273 if dpo_info[ "ready" ]:
274 recommended.append( "DPO" )
275 elif dpo_info[ "can_map" ]:
276 recommended.append( "DPO (needs mapping)" )
277
278 if grpo_info[ "ready" ]:
279 recommended.append( "GRPO" )
280 elif grpo_info[ "can_map" ]:
281 recommended.append( "GRPO (needs mapping)" )
282
283 if kto_info[ "ready" ]:
284 recommended.append( "KTO" )
285
286 # JSON output mode
287 if args.json_output:
288 result = {
289 "dataset" : args.dataset,
290 "config" : config_to_use,
291 "split" : args.split,
292 "total_examples" : total_examples,
293 "columns" : columns,
294 "features" : [{ "name" : f[ "name" ], "type" : f[ "type" ]} for f in features] if features else [],
295 "compatibility" : {
296 "SFT" : sft_info,
297 "DPO" : dpo_info,
298 "GRPO" : grpo_info,
299 "KTO" : kto_info,
300 },
301 "recommended_methods" : recommended,
302 }
303 print (json.dumps(result, indent = 2 ))
304 sys.exit( 0 )
305
306 # Human-readable output optimized for LLM parsing
307 print ( "=" * 80 )
308 print ( f "DATASET INSPECTION RESULTS" )
309 print ( "=" * 80 )
310
311 print ( f " \n Dataset: { args.dataset } " )
312 print ( f "Config: { config_to_use } " )
313 print ( f "Split: { args.split } " )
314 print ( f "Total examples: { total_examples } " )
315 print ( f "Samples fetched: { len (rows) } " )
316
317 print ( f " \n{ 'COLUMNS' : -< 80 } " )
318 if features:
319 for feature in features:
320 print ( f " { feature[ 'name' ] } : { feature[ 'type' ] } " )
321 else :
322 for col in columns:
323 print ( f " { col } : (type info not available)" )
324
325 print ( f " \n{ 'EXAMPLE DATA' : -< 80 } " )
326 example = first_row
327 for col in columns:
328 value = example.get(col)
329 display = format_value_preview(value, args.preview)
330 print ( f " \n{ col } :" )
331 print ( f " { display } " )
332
333 print ( f " \n{ 'TRAINING METHOD COMPATIBILITY' : -< 80 } " )
334
335 # SFT
336 print ( f " \n [SFT] { '✓ READY' if sft_info[ 'ready' ] else '✗ NEEDS MAPPING' } " )
337 if sft_info[ "ready" ]:
338 print ( f " Reason: Dataset has ' { sft_info[ 'reason' ] } ' field" )
339 print ( f " Action: Use directly with SFTTrainer" )
340 elif sft_info[ "possible_prompt" ]:
341 print ( f " Detected: prompt=' { sft_info[ 'possible_prompt' ] } ' response=' { sft_info[ 'possible_response' ] } '" )
342 print ( f " Action: Apply mapping code (see below)" )
343 else :
344 print ( f " Status: Cannot determine mapping - manual inspection needed" )
345
346 # DPO
347 print ( f " \n [DPO] { '✓ READY' if dpo_info[ 'ready' ] else '✗ NEEDS MAPPING' if dpo_info[ 'can_map' ] else '✗ INCOMPATIBLE' } " )
348 if dpo_info[ "ready" ]:
349 print ( f " Reason: Dataset has 'prompt', 'chosen', 'rejected' fields" )
350 print ( f " Action: Use directly with DPOTrainer" )
351 elif dpo_info[ "can_map" ]:
352 print ( f " Detected: prompt=' { dpo_info[ 'prompt_col' ] } ' chosen=' { dpo_info[ 'chosen_col' ] } ' rejected=' { dpo_info[ 'rejected_col' ] } '" )
353 print ( f " Action: Apply mapping code (see below)" )
354 else :
355 print ( f " Status: Missing required fields (prompt + chosen + rejected)" )
356
357 # GRPO
358 print ( f " \n [GRPO] { '✓ READY' if grpo_info[ 'ready' ] else '✗ NEEDS MAPPING' if grpo_info[ 'can_map' ] else '✗ INCOMPATIBLE' } " )
359 if grpo_info[ "ready" ]:
360 print ( f " Reason: Dataset has 'prompt' field" )
361 print ( f " Action: Use directly with GRPOTrainer" )
362 elif grpo_info[ "can_map" ]:
363 print ( f " Detected: prompt=' { grpo_info[ 'prompt_col' ] } '" )
364 print ( f " Action: Apply mapping code (see below)" )
365 else :
366 print ( f " Status: Missing prompt field" )
367
368 # KTO
369 print ( f " \n [KTO] { '✓ READY' if kto_info[ 'ready' ] else '✗ INCOMPATIBLE' } " )
370 if kto_info[ "ready" ]:
371 print ( f " Reason: Dataset has 'prompt', 'completion', 'label' fields" )
372 print ( f " Action: Use directly with KTOTrainer" )
373 else :
374 print ( f " Status: Missing required fields (prompt + completion + label)" )
375
376 # Mapping code
377 print ( f " \n{ 'MAPPING CODE (if needed)' : -< 80 } " )
378
379 mapping_needed = False
380
381 sft_mapping = generate_mapping_code( "SFT" , sft_info)
382 if sft_mapping:
383 print ( f " \n # For SFT Training:" )
384 print (sft_mapping)
385 mapping_needed = True
386
387 dpo_mapping = generate_mapping_code( "DPO" , dpo_info)
388 if dpo_mapping:
389 print ( f " \n # For DPO Training:" )
390 print (dpo_mapping)
391 mapping_needed = True
392
393 grpo_mapping = generate_mapping_code( "GRPO" , grpo_info)
394 if grpo_mapping:
395 print ( f " \n # For GRPO Training:" )
396 print (grpo_mapping)
397 mapping_needed = True
398
399 if not mapping_needed:
400 print ( " \n No mapping needed - dataset is ready for training!" )
401
402 print ( f " \n{ 'SUMMARY' : -< 80 } " )
403 print ( f "Recommended training methods: { ', ' .join(recommended) if recommended else 'None (dataset needs formatting)' } " )
404 print ( f " \n Note: Used Datasets Server API (instant, no download required)" )
405
406 print ( " \n " + "=" * 80 )
407 sys.exit( 0 )
408
409
410 if __name__ == "__main__" :
411 try :
412 main()
413 except KeyboardInterrupt :
414 sys.exit( 0 )
415 except Exception as e:
416 print ( f "ERROR: { e } " , file = sys.stderr)
417 sys.exit( 1 )