58 lines
1.9 KiB
Python
58 lines
1.9 KiB
Python
import re
|
|
from typing import Dict
|
|
import pandas as pd
|
|
|
|
def increment_count(d: Dict[str, int], category: str) -> None:
|
|
"""Increment the count of a category in the dictionary.
|
|
|
|
Args:
|
|
d (Dict[str, int]): Dictionary storing category counts.
|
|
category (str): The category to increment.
|
|
"""
|
|
d[category] = d.get(category, 0) + 1
|
|
|
|
def process_import_lines(file_path: str) -> Dict[str, int]:
|
|
"""Extract and categorize import lines from a Python file.
|
|
|
|
Args:
|
|
file_path (str): Path to the Python file.
|
|
|
|
Returns:
|
|
Dict[str, int]: Dictionary containing category counts.
|
|
"""
|
|
category_counts = {}
|
|
|
|
with open(file_path, "r") as f:
|
|
for line in f:
|
|
if "import" in line:
|
|
if line.startswith("from typing"):
|
|
continue
|
|
module_name = re.sub(r"^.*import ", "", line).strip()
|
|
if module_name.startswith("SP"):
|
|
increment_count(category_counts, "SentencePiece")
|
|
elif module_name.startswith("Pt"):
|
|
increment_count(category_counts, "PyTorch")
|
|
else:
|
|
increment_count(category_counts, "Misc")
|
|
|
|
return category_counts
|
|
|
|
def save_category_counts_to_csv(category_counts: Dict[str, int], output_file: str) -> None:
|
|
"""Save category counts as a sorted CSV file.
|
|
|
|
Args:
|
|
category_counts (Dict[str, int]): Dictionary containing category counts.
|
|
output_file (str): Path to the output CSV file.
|
|
"""
|
|
df = pd.DataFrame(list(category_counts.items()), columns=["Category", "Count"])
|
|
df = df.sort_values(by=["Count"], ascending=False)
|
|
df.to_csv(output_file, index=False)
|
|
|
|
def main() -> None:
|
|
"""Main function to process the import categories and save results."""
|
|
category_counts = process_import_lines("__init__.py")
|
|
save_category_counts_to_csv(category_counts, "node_count_breakdown.csv")
|
|
|
|
if __name__ == "__main__":
|
|
main()
|