countries.py (1963B)
1 import argparse 2 import collections 3 import pathlib 4 import re 5 6 import datasets 7 8 9 NUM_COUNTRIES = 12 10 11 12 def main(dataset_path: pathlib.Path) -> None: 13 split_size: dict[str, int] = {} 14 per_split: dict[str, dict[str, int]] = {} 15 for split_path in dataset_path.iterdir(): 16 if split_path.is_dir(): 17 split = datasets.load_from_disk(str(split_path)) 18 split_size[split_path.name] = len(split) 19 per_split[split_path.name] = collections.Counter(split["location_root_name"]) 20 21 max_rate_per_country: dict[str, float] = collections.Counter() 22 for split, countries in per_split.items(): 23 for country, cardinal in countries.items(): 24 rate: float = cardinal / split_size[split] 25 max_rate_per_country[country] = max(max_rate_per_country[country], rate) 26 27 selection: list[str] = [ country for country, _ in max_rate_per_country.most_common(NUM_COUNTRIES) ] 28 selection.sort(key=lambda country: per_split["train"][country], reverse=True) 29 30 for index, country in enumerate(selection): 31 short_country: str = re.sub(r" \(.*", "", country) 32 rates: dict[str, float] = { split: per_split[split][country] / split_size[split] for split in split_size } 33 print(r"\countryDistrib"\ 34 f"{{{index}}}"\ 35 f"{{{short_country}}}"\ 36 f"{{{rates['train']}}}"\ 37 f"{{{rates['validation']}}}"\ 38 f"{{{rates['test']}}}"\ 39 r"\\") 40 41 42 def parse_command_argument() -> argparse.Namespace: 43 """ Parse command line arguments. """ 44 parser = argparse.ArgumentParser(description="Generate country statistics for TikZ figure.") 45 parser.add_argument("dataset", type=pathlib.Path, help="Path to the target dataset splits.") 46 47 args: argparse.Namespace = parser.parse_args() 48 return args 49 50 51 if __name__ == "__main__": 52 args: argparse.Namespace = parse_command_argument() 53 main(dataset_path=args.dataset)