Setting the file. One moment.
Hf Benchmarks · Huggingface LLM Trainer · huggingface/skills · Skills Docs
Back to the top of the page 14.10
Unsloth
def collect_prefixed_tags
— line 213
This file
Number 14.14
Position 14 of 18
Type Python
Size 20 KB
Lines 659 scripts/ hf_benchmarks.py
Python · 659 lines · 20 KB
17 from __future__ import annotations
18
19 import argparse
20 import json
21 import os
22 import re
23 import sys
24 import textwrap
25 import urllib.error
26 import urllib.parse
27 import urllib.request
28 from typing import Any, Iterable
29
30
31 BASE_URL = "https://huggingface.co"
32 DEFAULT_TIMEOUT = 30
33
34 ALIASES : dict[ str , list[ str ]] = {
35 "ocr" : [
36 "ocr" ,
37 "olmocr" ,
38 "pdf" ,
39 "image-to-text" ,
40 "screen" ,
41 "screenspot" ,
42 "markdown" ,
43 "text recognition" ,
44 ],
45 "coding" : [
46 "code" ,
47 "coding" ,
48 "software engineering" ,
49 "programming" ,
50 "swe" ,
51 "terminal" ,
52 "patch" ,
53 "bug" ,
54 "cuda" ,
55 ],
56 "math" : [
57 "math" ,
58 "reasoning" ,
59 "gsm8k" ,
60 "mmlu" ,
61 "gpqa" ,
62 "aime" ,
63 "hmmt" ,
64 ],
65 "retrieval" : [
66 "retrieval" ,
67 "search" ,
68 "mteb" ,
69 "arguana" ,
70 "bright" ,
71 ],
72 "agents" : [
73 "agent" ,
74 "agents" ,
75 "terminal" ,
76 "screen" ,
77 "computer use" ,
78 "tool use" ,
79 ],
80 "asr" : [
81 "asr" ,
82 "speech" ,
83 "audio" ,
84 "transcribe" ,
85 "transcription" ,
86 ],
87 }
88
89
90 class HfApiError ( RuntimeError ):
91 pass
92
93
94 class FullHelpArgumentParser ( argparse . ArgumentParser ):
95 def __init__ (self, * args: Any, ** kwargs: Any) -> None :
96 super (). __init__ ( * args, ** kwargs)
97 self ._search_parser: argparse.ArgumentParser | None = None
98 self ._leaderboard_parser: argparse.ArgumentParser | None = None
99
100 def format_help (self) -> str :
101 text = super ().format_help()
102 extra_sections: list[ str ] = []
103
104 if self ._search_parser is not None :
105 extra_sections.append(
106 " \n search command options: \n "
107 + textwrap.indent( self ._search_parser.format_help().strip(), " " )
108 )
109
110 if self ._leaderboard_parser is not None :
111 extra_sections.append(
112 " \n leaderboard command options: \n "
113 + textwrap.indent( self ._leaderboard_parser.format_help().strip(), " " )
114 )
115
116 if extra_sections:
117 text += " \n " + " \n " .join(extra_sections) + " \n "
118 return text
119
120
121 def auth_headers () -> dict[ str , str ]:
122 token = os.getenv( "HF_TOKEN" )
123 return { "Authorization" : f "Bearer { token } " } if token else {}
124
125
126 def http_get_json (path: str , params: dict[ str , Any] | None = None ) -> Any:
127 url = f " { BASE_URL }{ path } "
128 if params:
129 pairs: list[tuple[ str , str ]] = []
130 for key, value in params.items():
131 if value is None :
132 continue
133 if isinstance (value, ( list , tuple )):
134 for item in value:
135 pairs.append((key, str (item)))
136 else :
137 pairs.append((key, str (value)))
138 if pairs:
139 url = f " { url } ? { urllib.parse.urlencode(pairs) } "
140
141 req = urllib.request.Request(url, headers = auth_headers())
142 try :
143 with urllib.request.urlopen(req, timeout = DEFAULT_TIMEOUT ) as resp:
144 return json.loads(resp.read().decode( "utf-8" ))
145 except urllib.error.HTTPError as exc:
146 body = exc.read().decode( "utf-8" , errors = "replace" )
147 raise HfApiError( f " { exc.code } { exc.reason } for { url } : { body[: 500 ] } " ) from exc
148 except urllib.error.URLError as exc:
149 raise HfApiError( f "Request failed for { url } : { exc } " ) from exc
150
151
152 def shorten (text: str , width: int ) -> str :
153 text = " " .join((text or "" ).split())
154 if len (text) <= width:
155 return text
156 return text[: max ( 0 , width - 1 )] + "…"
157
158
159 def first_text (value: Any) -> str :
160 if value is None :
161 return ""
162 if isinstance (value, str ):
163 return value
164 if isinstance (value, list ):
165 return " " .join(first_text(v) for v in value)
166 if isinstance (value, dict ):
167 return " " .join(first_text(v) for v in value.values())
168 return str (value)
169
170
171 def benchmark_catalog (limit: int = 500 ) -> list[dict[ str , Any]]:
172 data = http_get_json(
173 "/api/datasets" ,
174 params = { "filter" : "benchmark:official" , "limit" : limit, "full" : "true" },
175 )
176 if not isinstance (data, list ):
177 raise HfApiError( "Unexpected response while listing benchmark datasets" )
178 return data
179
180
181 def dataset_search_blob (dataset: dict[ str , Any]) -> str :
182 card = dataset.get( "cardData" ) or {}
183 parts = [
184 dataset.get( "id" , "" ),
185 dataset.get( "description" , "" ),
186 first_text(dataset.get( "tags" )),
187 first_text(card.get( "pretty_name" )),
188 first_text(card.get( "tags" )),
189 first_text(card.get( "task_categories" )),
190 first_text(card.get( "task_ids" )),
191 ]
192 return " " .join(parts).lower()
193
194
195 def dataset_search_fields (dataset: dict[ str , Any]) -> dict[ str , str ]:
196 card = dataset.get( "cardData" ) or {}
197 return {
198 "id" : first_text(dataset.get( "id" )).lower(),
199 "pretty_name" : first_text(card.get( "pretty_name" )).lower(),
200 "tags" : " " .join(
201 [
202 first_text(dataset.get( "tags" )),
203 first_text(card.get( "tags" )),
204 first_text(card.get( "task_categories" )),
205 first_text(card.get( "task_ids" )),
206 first_text(card.get( "modality" )),
207 ]
208 ).lower(),
209 "description" : first_text(dataset.get( "description" )).lower(),
210 }
211
212
213 def collect_prefixed_tags (dataset: dict[ str , Any], prefixes: Iterable[ str ]) -> list[ str ]:
214 prefixes = tuple (prefixes)
215 tags = dataset.get( "tags" ) or []
216 card = dataset.get( "cardData" ) or {}
217
218 out: list[ str ] = []
219 for tag in tags:
220 if isinstance (tag, str ) and tag.startswith(prefixes):
221 out.append(tag)
222
223 for key, prefix in (
224 ( "task_categories" , "task_categories:" ),
225 ( "task_ids" , "task_ids:" ),
226 ( "modality" , "modality:" ),
227 ):
228 values = card.get(key)
229 if isinstance (values, list ):
230 for value in values:
231 full_tag = f " { prefix }{ value } "
232 if full_tag.startswith(prefixes):
233 out.append(full_tag)
234
235 deduped: list[ str ] = []
236 seen: set[ str ] = set ()
237 for tag in out:
238 if tag not in seen:
239 deduped.append(tag)
240 seen.add(tag)
241 return deduped
242
243
244 def expand_aliases (aliases: list[ str ]) -> dict[ str , list[ str ]]:
245 expanded: dict[ str , list[ str ]] = {}
246 for alias in aliases:
247 terms = ALIASES .get(alias.lower(), [alias])
248 expanded[alias] = terms
249 return expanded
250
251
252 def matches_term (blob: str , term: str ) -> bool :
253 candidate = term.lower().strip()
254 if not candidate:
255 return False
256 if re.fullmatch( r " [ a-z0-9_ ] + " , candidate):
257 return re.search( rf "(?<![a-z0-9_]) { re.escape(candidate) } (?![a-z0-9_])" , blob) is not None
258 return candidate in blob
259
260
261 def score_dataset (
262 dataset: dict[ str , Any],
263 queries: list[ str ],
264 aliases: dict[ str , list[ str ]],
265 tasks: list[ str ],
266 modalities: list[ str ],
267 ) -> dict[ str , Any]:
268 blob = dataset_search_blob(dataset)
269 fields = dataset_search_fields(dataset)
270 task_tags = collect_prefixed_tags(dataset, [ "task_categories:" , "task_ids:" ])
271 modality_tags = collect_prefixed_tags(dataset, [ "modality:" ])
272
273 score = 0
274 reasons: list[ str ] = []
275
276 for query in queries:
277 q = query.lower().strip()
278 if q and any (matches_term(value, q) for value in fields.values()):
279 score += 3
280 reasons.append( f "query: { query } " )
281
282 for alias_name, terms in aliases.items():
283 matched_terms: list[ str ] = []
284 alias_score = 0
285 for term in terms:
286 strong_match = any (
287 matches_term(fields[field_name], term)
288 for field_name in ( "id" , "pretty_name" , "tags" )
289 )
290 desc_match = matches_term(fields[ "description" ], term)
291 if strong_match:
292 alias_score += 2
293 matched_terms.append(term)
294 elif desc_match:
295 alias_score += 1
296 matched_terms.append(term)
297 if matched_terms:
298 score += alias_score
299 reasons.append( f "alias: { alias_name } =" + "," .join(matched_terms[: 5 ]))
300
301 lower_task_tags = [t.lower() for t in task_tags]
302 for task in tasks:
303 task = task.lower().strip()
304 if not task:
305 continue
306 exact = [
307 tag
308 for tag in lower_task_tags
309 if tag == f "task_categories: { task } " or tag == f "task_ids: { task } "
310 ]
311 fuzzy = matches_term(blob, task)
312 if exact:
313 score += 5
314 reasons.append( f "task: { task } " )
315 elif fuzzy:
316 score += 2
317 reasons.append( f "task~: { task } " )
318
319 lower_modality_tags = [m.lower() for m in modality_tags]
320 for modality in modalities:
321 modality = modality.lower().strip()
322 if not modality:
323 continue
324 if f "modality: { modality } " in lower_modality_tags:
325 score += 4
326 reasons.append( f "modality: { modality } " )
327
328 return {
329 "dataset_id" : dataset.get( "id" ),
330 "score" : score,
331 "reasons" : reasons,
332 "task_tags" : task_tags,
333 "modality_tags" : modality_tags,
334 "benchmark_tags" : collect_prefixed_tags(dataset, [ "benchmark:" ]),
335 "pretty_name" : (dataset.get( "cardData" ) or {}).get( "pretty_name" ),
336 "downloads" : dataset.get( "downloads" ),
337 "description" : " " .join((dataset.get( "description" ) or "" ).split()),
338 }
339
340
341 def search_benchmarks (
342 queries: list[ str ],
343 aliases: list[ str ],
344 tasks: list[ str ],
345 modalities: list[ str ],
346 limit: int ,
347 ) -> list[dict[ str , Any]]:
348 datasets = benchmark_catalog( limit = 500 )
349 alias_map = expand_aliases(aliases)
350
351 results = [score_dataset(ds, queries, alias_map, tasks, modalities) for ds in datasets]
352
353 active_filters = bool (queries or aliases or tasks or modalities)
354 if active_filters:
355 results = [row for row in results if row[ "score" ] >= 2 ]
356
357 results.sort(
358 key =lambda row: (
359 - row[ "score" ],
360 - (row[ "downloads" ] or 0 ),
361 row[ "dataset_id" ] or "" ,
362 )
363 )
364 return results[:limit]
365
366
367 def parse_repo_id (repo_id: str ) -> tuple[ str , str ]:
368 if "/" not in repo_id:
369 raise ValueError ( f "Expected <namespace>/<repo>, got: { repo_id } " )
370 namespace, repo = repo_id.split( "/" , 1 )
371 return namespace, repo
372
373
374 def get_leaderboard (repo_id: str , task_id: str | None = None ) -> list[dict[ str , Any]]:
375 namespace, repo = parse_repo_id(repo_id)
376 params = { "task_id" : task_id} if task_id else None
377 data = http_get_json( f "/api/datasets/ { namespace } / { repo } /leaderboard" , params = params)
378 if not isinstance (data, list ):
379 raise HfApiError( f "Unexpected leaderboard response for { repo_id } " )
380
381 normalized: list[dict[ str , Any]] = []
382 for row in data:
383 source = row.get( "source" ) or {}
384 normalized.append(
385 {
386 "dataset_id" : repo_id,
387 "task_id" : task_id,
388 "rank" : row.get( "rank" ),
389 "model_id" : row.get( "modelId" ),
390 "value" : row.get( "value" ),
391 "verified" : row.get( "verified" ),
392 "lower_is_better" : row.get( "lower_is_better" ),
393 "filename" : row.get( "filename" ),
394 "notes" : row.get( "notes" ),
395 "pull_request" : row.get( "pullRequest" ),
396 "source_name" : source.get( "name" ),
397 "source_url" : source.get( "url" ),
398 "source_is_external" : source.get( "isExternal" ),
399 }
400 )
401 return normalized
402
403
404 def read_repo_ids_from_stdin () -> list[ str ]:
405 if sys.stdin.isatty():
406 return []
407
408 repo_ids: list[ str ] = []
409 for raw_line in sys.stdin:
410 line = raw_line.strip()
411 if not line:
412 continue
413 if line.startswith( "{" ):
414 try :
415 obj = json.loads(line)
416 except json.JSONDecodeError:
417 continue
418 candidate = obj.get( "dataset_id" ) or obj.get( "id" )
419 if isinstance (candidate, str ) and "/" in candidate:
420 repo_ids.append(candidate)
421 continue
422 if "/" in line:
423 repo_ids.append(line)
424 return repo_ids
425
426
427 def print_json (data: Any) -> None :
428 json.dump(data, sys.stdout, indent = 2 , ensure_ascii = False )
429 sys.stdout.write( " \n " )
430
431
432 def print_ndjson (rows: list[dict[ str , Any]]) -> None :
433 for row in rows:
434 sys.stdout.write(json.dumps(row, ensure_ascii = False ) + " \n " )
435
436
437 def print_search_table (rows: list[dict[ str , Any]]) -> None :
438 if not rows:
439 print ( "No benchmark datasets matched." )
440 return
441
442 headers = [ "dataset_id" , "score" , "modalities" , "tasks" , "reasons" , "description" ]
443 widths = [ 34 , 5 , 18 , 24 , 30 , 68 ]
444 print ( " " .join(h.ljust(w) for h, w in zip (headers, widths)))
445 print ( " " .join( "-" * w for w in widths))
446 for row in rows:
447 values = [
448 shorten(row.get( "dataset_id" ) or "" , widths[ 0 ]),
449 str (row.get( "score" , "" )),
450 shorten( ", " .join(row.get( "modality_tags" ) or []), widths[ 2 ]),
451 shorten( ", " .join(row.get( "task_tags" ) or []), widths[ 3 ]),
452 shorten( ", " .join(row.get( "reasons" ) or []), widths[ 4 ]),
453 shorten(row.get( "description" ) or "" , widths[ 5 ]),
454 ]
455 print ( " " .join(v.ljust(w) for v, w in zip (values, widths)))
456
457
458 def print_leaderboard_table (rows: list[dict[ str , Any]]) -> None :
459 if not rows:
460 print ( "No leaderboard rows returned." )
461 return
462
463 headers = [ "dataset_id" , "rank" , "model_id" , "value" , "verified" , "source" ]
464 widths = [ 30 , 5 , 38 , 10 , 8 , 28 ]
465 print ( " " .join(h.ljust(w) for h, w in zip (headers, widths)))
466 print ( " " .join( "-" * w for w in widths))
467 for row in rows:
468 values = [
469 shorten( str (row.get( "dataset_id" ) or "" ), widths[ 0 ]),
470 str (row.get( "rank" ) or "" ),
471 shorten( str (row.get( "model_id" ) or "" ), widths[ 2 ]),
472 shorten( str (row.get( "value" ) or "" ), widths[ 3 ]),
473 str (row.get( "verified" )),
474 shorten( str (row.get( "source_name" ) or "" ), widths[ 5 ]),
475 ]
476 print ( " " .join(v.ljust(w) for v, w in zip (values, widths)))
477
478
479 def build_parser () -> argparse.ArgumentParser:
480 parser = FullHelpArgumentParser(
481 prog = "hf_benchmarks.py" ,
482 formatter_class = argparse.RawDescriptionHelpFormatter,
483 description = textwrap.dedent(
484 """
485 Search benchmark datasets and fetch leaderboard results from the Hugging Face Hub.
486
487 Workflow ideas:
488 1) Discover candidate benchmarks:
489 hf_benchmarks.py search --alias ocr
490 hf_benchmarks.py search --alias coding
491 hf_benchmarks.py search --task image-to-text --modality document
492
493 2) Inspect a leaderboard:
494 hf_benchmarks.py leaderboard allenai/olmOCR-bench --top 10
495
496 3) Chain search -> leaderboard:
497 hf_benchmarks.py search --alias coding --format ndjson \\
498 | hf_benchmarks.py leaderboard --stdin --top 5 --format table
499 """
500 ),
501 )
502
503 subparsers = parser.add_subparsers( dest = "command" , required = True )
504
505 search_parser = subparsers.add_parser(
506 "search" ,
507 help = "Search benchmark datasets by query, alias, task, and modality" ,
508 )
509 search_parser.add_argument(
510 "--query" ,
511 action = "append" ,
512 default = [],
513 help = "Free-text query to match against benchmark dataset metadata. Repeatable." ,
514 )
515 search_parser.add_argument(
516 "--alias" ,
517 action = "append" ,
518 default = [],
519 help = (
520 "Convenience alias for common benchmark domains. Known aliases: "
521 + ", " .join( sorted ( ALIASES ))
522 + ". Repeatable."
523 ),
524 )
525 search_parser.add_argument(
526 "--task" ,
527 action = "append" ,
528 default = [],
529 help = "Task to match, e.g. text-generation, image-to-text, question-answering. Repeatable." ,
530 )
531 search_parser.add_argument(
532 "--modality" ,
533 action = "append" ,
534 default = [],
535 help = "Modality to match, e.g. text, image, document, audio. Repeatable." ,
536 )
537 search_parser.add_argument(
538 "--limit" ,
539 type = int ,
540 default = 20 ,
541 help = "Maximum number of rows to print (default: 20)." ,
542 )
543 search_parser.add_argument(
544 "--format" ,
545 choices = [ "table" , "json" , "ndjson" ],
546 default = "table" ,
547 help = "Output format (default: table)." ,
548 )
549
550 leaderboard_parser = subparsers.add_parser(
551 "leaderboard" ,
552 help = "Fetch normalized leaderboard rows for one or more benchmark datasets" ,
553 )
554 leaderboard_parser.add_argument(
555 "datasets" ,
556 nargs = "*" ,
557 help = "Dataset repo ids (<namespace>/<repo>). Can also be supplied via stdin with --stdin." ,
558 )
559 leaderboard_parser.add_argument(
560 "--stdin" ,
561 action = "store_true" ,
562 help = "Read dataset ids from stdin. Accepts plain repo ids or NDJSON with dataset_id/id fields." ,
563 )
564 leaderboard_parser.add_argument(
565 "--task-id" ,
566 default = None ,
567 help = "Optional leaderboard task_id query parameter." ,
568 )
569 leaderboard_parser.add_argument(
570 "--top" ,
571 type = int ,
572 default = None ,
573 help = "Only keep the top N results per leaderboard." ,
574 )
575 leaderboard_parser.add_argument(
576 "--format" ,
577 choices = [ "table" , "json" , "ndjson" ],
578 default = "table" ,
579 help = "Output format (default: table)." ,
580 )
581
582 parser._search_parser = search_parser
583 parser._leaderboard_parser = leaderboard_parser
584
585 return parser
586
587
588 def run_search (args: argparse.Namespace) -> int :
589 rows = search_benchmarks(
590 queries = args.query,
591 aliases = args.alias,
592 tasks = args.task,
593 modalities = args.modality,
594 limit = args.limit,
595 )
596
597 if args.format == "json" :
598 print_json(rows)
599 elif args.format == "ndjson" :
600 print_ndjson(rows)
601 else :
602 print_search_table(rows)
603 return 0
604
605
606 def run_leaderboard (args: argparse.Namespace) -> int :
607 repo_ids = list (args.datasets)
608 if args.stdin:
609 repo_ids.extend(read_repo_ids_from_stdin())
610
611 deduped: list[ str ] = []
612 seen: set[ str ] = set ()
613 for repo_id in repo_ids:
614 if repo_id not in seen:
615 deduped.append(repo_id)
616 seen.add(repo_id)
617 repo_ids = deduped
618
619 if not repo_ids:
620 print ( "Error: provide dataset ids or use --stdin." , file = sys.stderr)
621 return 2
622
623 rows: list[dict[ str , Any]] = []
624 for repo_id in repo_ids:
625 dataset_rows = get_leaderboard(repo_id, task_id = args.task_id)
626 if args.top is not None :
627 dataset_rows = dataset_rows[: args.top]
628 rows.extend(dataset_rows)
629
630 if args.format == "json" :
631 print_json(rows)
632 elif args.format == "ndjson" :
633 print_ndjson(rows)
634 else :
635 print_leaderboard_table(rows)
636 return 0
637
638
639 def main () -> int :
640 parser = build_parser()
641 args = parser.parse_args()
642
643 try :
644 if args.command == "search" :
645 return run_search(args)
646 if args.command == "leaderboard" :
647 return run_leaderboard(args)
648 parser.error( f "Unknown command: { args.command } " )
649 return 2
650 except HfApiError as exc:
651 print ( f "Error: { exc } " , file = sys.stderr)
652 return 1
653 except ValueError as exc:
654 print ( f "Error: { exc } " , file = sys.stderr)
655 return 1
656
657
658 if __name__ == "__main__" :
659 raise SystemExit (main())