transformers/utils/create_dependency_mapping.py
Arthur a7f5479b45
fix modular order (#35297)
* fix modular ordre

* fix

* style
2024-12-17 08:05:35 +01:00

78 lines
2.6 KiB
Python

import ast
from collections import defaultdict
# Function to perform topological sorting
def topological_sort(dependencies):
new_dependencies = {}
graph = defaultdict(list)
for node, deps in dependencies.items():
for dep in deps:
if "example" not in node and "auto" not in dep:
graph[dep.split(".")[-2]].append(node.split("/")[-2])
new_dependencies[node.split("/")[-2]] = node
# Create a graph and in-degree count for each node
def filter_one_by_one(filtered_list, reverse):
if len(reverse) == 0:
return filtered_list
graph = defaultdict(list)
# Build the graph
for node, deps in reverse.items():
for dep in deps:
graph[dep].append(node)
base_modules = set(reverse.keys()) - set(graph.keys())
if base_modules == reverse.keys():
# we are at the end
return filtered_list + list(graph.keys())
to_add = []
for k in graph.keys():
if len(graph[k]) == 1 and graph[k][0] in base_modules:
if graph[k][0] in reverse:
del reverse[graph[k][0]]
if k not in filtered_list:
to_add += [k]
for k in base_modules:
if k not in filtered_list:
to_add += [k]
filtered_list += list(to_add)
return filter_one_by_one(filtered_list, reverse)
final_order = filter_one_by_one([], graph)
return [new_dependencies.get(k) for k in final_order if k in new_dependencies]
# Function to extract class and import info from a file
def extract_classes_and_imports(file_path):
with open(file_path, "r") as file:
tree = ast.parse(file.read(), filename=file_path)
imports = set()
for node in ast.walk(tree):
if isinstance(node, (ast.Import, ast.ImportFrom)):
module = node.module if isinstance(node, ast.ImportFrom) else None
if module and (".modeling_" in module):
imports.add(module)
return imports
# Function to map dependencies between classes
def map_dependencies(py_files):
dependencies = defaultdict(set)
# First pass: Extract all classes and map to files
for file_path in py_files:
# dependencies[file_path].add(None)
class_to_file = extract_classes_and_imports(file_path)
for module in class_to_file:
dependencies[file_path].add(module)
return dependencies
def find_priority_list(py_files):
dependencies = map_dependencies(py_files)
ordered_classes = topological_sort(dependencies)
return ordered_classes