import argparse import difflib import glob import logging import multiprocessing import os import shutil import subprocess from functools import partial from create_dependency_mapping import find_priority_list # Console for rich printing from modular_model_converter import convert_modular_file, run_ruff from rich.console import Console from rich.syntax import Syntax CHECKER_CONFIG = { "name": "modular_conversion", "label": "Modular file conversions", # Globs the modular sources; also reads generated modeling_*.py at runtime for diffing. "cache_globs": ["src/transformers/models/**/modular_*.py", "src/transformers/models/**/modeling_*.py"], "check_args": [], "fix_args": ["--fix_and_overwrite"], } logging.basicConfig() logging.getLogger().setLevel(logging.ERROR) console = Console() BACKUP_EXT = ".modular_backup" def process_file( modular_file_path, generated_modeling_content, file_type="modeling_", show_diff=True, ): file_name_prefix = file_type.split(".*")[0] file_name_suffix = file_type.split(".*")[-1] if ".*" in file_type else "" file_path = modular_file_path.replace("modular_", f"{file_name_prefix}_").replace(".py", f"{file_name_suffix}.py") # Read the actual modeling file with open(file_path, "r", encoding="utf-8") as modeling_file: content = modeling_file.read() diff = difflib.unified_diff( generated_modeling_content[file_type].splitlines(), content.splitlines(), fromfile=f"{file_path}_generated", tofile=f"{file_path}", lineterm="", ) diff_list = list(diff) # Check for differences if diff_list: # first save the copy of the original file, to be able to restore it later shutil.copy(file_path, file_path + BACKUP_EXT) # we always save the generated content, to be able to update dependant files with open(file_path, "w", encoding="utf-8", newline="\n") as modeling_file: modeling_file.write(generated_modeling_content[file_type]) if not show_diff: console.print(f"[bold blue]Overwritten {file_path} with the generated content.[/bold blue]") if show_diff: console.print(f"\n[bold red]Differences found between the generated code and {file_path}:[/bold red]\n") diff_text = "\n".join(diff_list) syntax = Syntax(diff_text, "diff", theme="ansi_dark", line_numbers=True) console.print(syntax) return 1 else: return 0 def convert_and_run_ruff(modular_file_path: str) -> dict[str, str]: """From a modular file, convert it and return all the contents of the file as string. We need this function, because `ruff` needs the final filename to apply all rules correctly, so to get the output as a string, we need to save a temporary file with similar name, run ruff, and re-read the temporary file""" # Generate the expected modeling content generated_modeling_content = convert_modular_file(modular_file_path) # Temporary save the files with similar names to run `ruff` correctly, then re-read the result after linting/formatting for file_type in generated_modeling_content: file_name_prefix = file_type.split(".*")[0] file_name_suffix = file_type.split(".*")[-1] if ".*" in file_type else "" temp_file_name = modular_file_path.replace("modular_", f"{file_name_prefix}_").replace( ".py", f"_temp_pattern__{file_name_suffix}.py" ) # Write the file only temporarily with open(temp_file_name, "w") as f: f.write(generated_modeling_content[file_type]) # Run ruff on the new file (with similar name pattern as the original one) run_ruff(temp_file_name) with open(temp_file_name, "r") as f: generated_modeling_content[file_type] = f.read() # delete file os.remove(temp_file_name) return generated_modeling_content def compare_files(modular_file_path, show_diff=True): # Generate the expected modeling content generated_modeling_content = convert_and_run_ruff(modular_file_path) diff = 0 for file_type in generated_modeling_content: diff += process_file(modular_file_path, generated_modeling_content, file_type, show_diff) return diff # Changes to any of these files can alter the generated output for every modular model, # so touching them must force a full re-check (see `converter_changed_in_diff`). CONVERTER_FILES = { "utils/modular_model_converter.py", "utils/create_dependency_mapping.py", } def _get_modified_files(): fork_point_sha = subprocess.check_output("git merge-base main HEAD".split()).decode("utf-8") return ( subprocess.check_output(f"git diff --diff-filter=d --name-only {fork_point_sha}".split()) .decode("utf-8") .split() ) def get_models_in_diff(): """ Finds all models that have been modified in the diff. Returns: A set containing the names of the models that have been modified (e.g. {'llama', 'whisper'}). """ modified_files = _get_modified_files() # Matches both modelling files and tests relevant_modified_files = [x for x in modified_files if "/models/" in x and x.endswith(".py")] model_names = set() for file_path in relevant_modified_files: model_name = file_path.split("/")[-2] model_names.add(model_name) return model_names def converter_changed_in_diff(): """Whether the diff touches a file that can change conversion output for every model.""" return any(f in CONVERTER_FILES for f in _get_modified_files()) def guaranteed_no_diff(modular_file_path, dependencies, models_in_diff): """ Returns whether it is guaranteed to have no differences between the modular file and the modeling file. Model is in the diff -> not guaranteed to have no differences Dependency is in the diff -> not guaranteed to have no differences Otherwise -> guaranteed to have no differences Args: modular_file_path: The path to the modular file. dependencies: A dictionary containing the dependencies of each modular file. models_in_diff: A set containing the names of the models that have been modified. Returns: A boolean indicating whether the model (code and tests) is guaranteed to have no differences. """ model_name = modular_file_path.rsplit("modular_", 1)[1].replace(".py", "") if model_name in models_in_diff: return False for dep in dependencies[modular_file_path]: # two possible patterns: `transformers.models.model_name.(...)` or `model_name.(...)` dependency_model_name = dep.split(".")[-2] if dependency_model_name in models_in_diff: return False return True if __name__ == "__main__": parser = argparse.ArgumentParser(description="Compare modular_xxx.py files with modeling_xxx.py files.") parser.add_argument( "--files", default=["all"], type=str, nargs="+", help="List of modular_xxx.py files to compare." ) parser.add_argument( "--fix_and_overwrite", action="store_true", help="Overwrite the modeling_xxx.py file if differences are found." ) parser.add_argument("--check_all", action="store_true", help="Check all files, not just the ones in the diff.") parser.add_argument( "--num_workers", default=-1, type=int, help="The number of workers to run. Default is -1, which means the number of CPU cores.", ) args = parser.parse_args() if args.files == ["all"]: args.files = glob.glob("src/transformers/models/**/modular_*.py", recursive=True) if args.num_workers == -1: args.num_workers = multiprocessing.cpu_count() # Assuming there is a topological sort on the dependency mapping: if the file being checked and its dependencies # are not in the diff, then there it is guaranteed to have no differences. If no models are in the diff, then this # script will do nothing. current_branch = subprocess.check_output(["git", "branch", "--show-current"], text=True).strip() if current_branch == "main": console.print( "[bold red]You are developing on the main branch. We cannot identify the list of changed files and will have to check all files. This may take a while.[/bold red]" ) models_in_diff = {file_path.split("/")[-2] for file_path in args.files} elif converter_changed_in_diff(): # The converter (or its dependency-mapping helper) is in the diff: its output can shift # for any model, so restrict-by-diff would miss regressions. Force a full check. console.print("[bold yellow]Converter change detected in diff; checking all modular files.[/bold yellow]") args.check_all = True models_in_diff = {file_path.split("/")[-2] for file_path in args.files} else: models_in_diff = get_models_in_diff() if not models_in_diff and not args.check_all: exit(0) non_matching_files = [] ordered_files, dependencies = find_priority_list(args.files) flat_ordered_files = [item for sublist in ordered_files for item in sublist] # ordered_files is a *sorted* list of lists of filepaths # - files from the first list do NOT depend on other files # - files in the second list depend on files from the first list # - files in the third list depend on files from the second and (optionally) the first list # - ... and so on # files (models) within the same list are *independent* of each other; # we start applying modular conversion to each list in parallel, starting from the first list pool = None pool_size = 0 try: for dependency_level_files in ordered_files: # Filter files guaranteed no diff files_to_check = [] for file_path in dependency_level_files: if args.check_all or not guaranteed_no_diff(file_path, dependencies, models_in_diff): files_to_check.append(file_path) if not files_to_check: continue required_pool_size = min(args.num_workers, len(files_to_check)) if pool is None or required_pool_size > pool_size * 4: # Only create a new pool if we don't have one yet or the current one # is too small. Creating new pools is expensive due to the imports in # the workers. if pool is not None: pool.terminate() pool.join() pool_size = required_pool_size pool = multiprocessing.Pool(processes=pool_size) try: is_changed_flags = pool.map( partial(compare_files, show_diff=not args.fix_and_overwrite), files_to_check, ) except Exception as e: console.print(f"[bold red]Failed to convert one or more files in batch: {files_to_check}[/bold red]") console.print(f"[bold red]Error: {e}[/bold red]") # Try to process files individually to identify which one failed is_changed_flags = [] for file_path in files_to_check: try: result = compare_files(file_path, show_diff=not args.fix_and_overwrite) is_changed_flags.append(result) except Exception as individual_error: console.print(f"[bold red]Failed to convert {file_path}: {individual_error}[/bold red]") is_changed_flags.append(0) # Mark as no change to continue processing # Collect changed files and their original paths for is_changed, file_path in zip(is_changed_flags, files_to_check): if is_changed: non_matching_files.append(file_path) # Update changed models, after each round of conversions # (save model folder name) models_in_diff.add(file_path.split("/")[-2]) finally: if pool is not None: pool.terminate() pool.join() # Restore overwritten files by modular (if needed) backup_files = glob.glob("**/*" + BACKUP_EXT, recursive=True) for backup_file_path in backup_files: overwritten_path = backup_file_path.replace(BACKUP_EXT, "") if not args.fix_and_overwrite and os.path.exists(overwritten_path): shutil.copy(backup_file_path, overwritten_path) os.remove(backup_file_path) if non_matching_files and not args.fix_and_overwrite: diff_models = set(file_path.split("/")[-2] for file_path in non_matching_files) # noqa models_str = "\n - " + "\n - ".join(sorted(diff_models)) raise ValueError(f"Some diff and their modeling code did not match. Models in diff:{models_str}")