utils.update_config

 1import _ast
 2import ast
 3
 4config_file_path = "./config/config.py"
 5
 6#: Workflow selection for demo notebook
 7UPDATED_WORKFLOWS = {
 8    "embedding_selection": {
 9        "mode": "compute",
10        "parameters": {
11            "compute_representativeness": 0.99,
12            "compute_unique_images_greedy": 0.01,
13            "compute_unique_images_deterministic": 0.99,
14            "compute_similar_images": 0.03,
15            "neighbour_count": 3,
16        },
17        "embedding_models": [
18            "detection-transformer-torch",
19            "zero-shot-detection-transformer-torch",
20            "clip-vit-base32-torch",
21        ],
22    },
23}
24
25
26class ConfigVisitor(ast.NodeTransformer):
27    """AST visitor that transforms specific configuration variable assignments in the source code."""
28
29    def visit_Assign(self, node):
30        """Processes AST nodes related to variable assignments, modifying specific configuration variables."""
31        if isinstance(node.targets[0], _ast.Name):
32            if node.targets[0].id == "SELECTED_WORKFLOW":
33                node.value = ast.Constant(value=["embedding_selection"])
34            elif node.targets[0].id == "WORKFLOWS":
35                node.value = ast.Constant(value=UPDATED_WORKFLOWS)
36            elif node.targets[0].id == "WANDB_ACTIVE":
37                node.value = ast.Constant(value=False)
38            elif node.targets[0].id == "HF_DO_UPLOAD":
39                node.value = ast.Constant(value=False)
40            elif node.targets[0].id == "V51_REMOTE":
41                node.value = ast.Constant(value=False)
42        return node
43
44
45if __name__ == "__main__":
46    # Transform the AST
47    transformer = ConfigVisitor()
48
49    with open(config_file_path, "r") as file:
50        content = file.read()
51
52    parsed_ast = ast.parse(content)
53    updated_ast = transformer.visit(parsed_ast)
54
55    # Convert AST back to source code
56
57    updated_content = ast.unparse(updated_ast)
58
59    with open(config_file_path, "w") as file:
60        file.write(updated_content)
61
62    print("Config file updated successfully.")