File size: 4,082 Bytes
de4ade4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 |
# Copyright 2022 MosaicML LLM Foundry authors
# SPDX-License-Identifier: Apache-2.0
import ast
import importlib
import os
from typing import List, Optional
__all__ = ['edit_files_for_hf_compatibility']
class DeleteSpecificNodes(ast.NodeTransformer):
def __init__(self, nodes_to_remove: List[ast.AST]):
self.nodes_to_remove = nodes_to_remove
def visit(self, node: ast.AST) -> Optional[ast.AST]:
if node in self.nodes_to_remove:
return None
return super().visit(node)
def convert_to_relative_import(
module_name: str, original_parent_module_name: Optional[str]) -> str:
parts = module_name.split('.')
if parts[-1] == original_parent_module_name:
return '.'
return '.' + parts[-1]
def find_module_file(module_name: str) -> str:
if not module_name:
raise ValueError(f'Invalid input: {module_name=}')
module = importlib.import_module(module_name)
module_file = module.__file__
if module_file is None:
raise ValueError(f'Could not find file for module: {module_name}')
return module_file
def process_file(file_path: str, folder_path: str) -> List[str]:
with open(file_path, 'r') as f:
source = f.read()
parent_module_name = None
if os.path.basename(file_path) == '__init__.py':
parent_module_name = os.path.basename(os.path.dirname(file_path))
tree = ast.parse(source)
new_files_to_process = []
nodes_to_remove = []
for node in ast.walk(tree):
# convert any llmfoundry imports into relative imports
if isinstance(
node, ast.ImportFrom
) and node.module is not None and node.module.startswith('llmfoundry'):
module_path = find_module_file(node.module)
node.module = convert_to_relative_import(node.module,
parent_module_name)
# recursively process any llmfoundry files
new_files_to_process.append(module_path)
# remove any imports from composer or omegaconf
elif isinstance(node, ast.ImportFrom) and node.module is not None and (
node.module.startswith('composer') or
node.module.startswith('omegaconf')):
nodes_to_remove.append(node)
# remove the Composer* class
elif isinstance(node,
ast.ClassDef) and node.name.startswith('Composer'):
nodes_to_remove.append(node)
# remove the __all__ declaration in any __init__.py files, whose enclosing module
# will be converted to a single file of the same name
elif isinstance(node,
ast.Assign) and len(node.targets) == 1 and isinstance(
node.targets[0],
ast.Name) and node.targets[0].id == '__all__':
nodes_to_remove.append(node)
transformer = DeleteSpecificNodes(nodes_to_remove)
new_tree = transformer.visit(tree)
new_filename = os.path.basename(file_path)
# special case for __init__.py to mimic the original submodule
if new_filename == '__init__.py':
new_filename = file_path.split('/')[-2] + '.py'
new_file_path = os.path.join(folder_path, new_filename)
with open(new_file_path, 'w') as f:
assert new_tree is not None
f.write(ast.unparse(new_tree))
return new_files_to_process
def edit_files_for_hf_compatibility(folder: str) -> None:
files_to_process = [
os.path.join(folder, filename)
for filename in os.listdir(folder)
if filename.endswith('.py')
]
files_processed_and_queued = set(files_to_process)
while len(files_to_process) > 0:
to_process = files_to_process.pop()
if os.path.isfile(to_process) and to_process.endswith('.py'):
to_add = process_file(to_process, folder)
for file in to_add:
if file not in files_processed_and_queued:
files_to_process.append(file)
files_processed_and_queued.add(file)
|