diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml new file mode 100644 index 0000000..9a04f17 --- /dev/null +++ b/.github/workflows/codeql.yml @@ -0,0 +1,100 @@ +# For most projects, this workflow file will not need changing; you simply need +# to commit it to your repository. +# +# You may wish to alter this file to override the set of languages analyzed, +# or to provide custom queries or build logic. +# +# ******** NOTE ******** +# We have attempted to detect the languages in your repository. Please check +# the `language` matrix defined below to confirm you have the correct set of +# supported CodeQL languages. +# +name: "CodeQL Advanced" + +on: + push: + branches: [ "main", "deploy" ] + pull_request: + branches: [ "main", "deploy" ] + workflow_dispatch: + +jobs: + analyze: + name: Analyze (${{ matrix.language }}) + # Runner size impacts CodeQL analysis time. To learn more, please see: + # - https://gh.io/recommended-hardware-resources-for-running-codeql + # - https://gh.io/supported-runners-and-hardware-resources + # - https://gh.io/using-larger-runners (GitHub.com only) + # Consider using larger runners or machines with greater resources for possible analysis time improvements. + runs-on: ${{ (matrix.language == 'swift' && 'macos-latest') || 'ubuntu-latest' }} + permissions: + # required for all workflows + security-events: write + + # required to fetch internal or private CodeQL packs + packages: read + + # only required for workflows in private repositories + actions: read + contents: read + + strategy: + fail-fast: false + matrix: + include: + - language: javascript-typescript + build-mode: none + - language: python + build-mode: none + # CodeQL supports the following values keywords for 'language': 'actions', 'c-cpp', 'csharp', 'go', 'java-kotlin', 'javascript-typescript', 'python', 'ruby', 'rust', 'swift' + # Use `c-cpp` to analyze code written in C, C++ or both + # Use 'java-kotlin' to analyze code written in Java, Kotlin or both + # Use 'javascript-typescript' to analyze code written in JavaScript, TypeScript or both + # To learn more about changing the languages that are analyzed or customizing the build mode for your analysis, + # see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/customizing-your-advanced-setup-for-code-scanning. + # If you are analyzing a compiled language, you can modify the 'build-mode' for that language to customize how + # your codebase is analyzed, see https://docs.github.com/en/code-security/code-scanning/creating-an-advanced-setup-for-code-scanning/codeql-code-scanning-for-compiled-languages + steps: + - name: Checkout repository + uses: actions/checkout@v4 + + # Add any setup steps before running the `github/codeql-action/init` action. + # This includes steps like installing compilers or runtimes (`actions/setup-node` + # or others). This is typically only required for manual builds. + # - name: Setup runtime (example) + # uses: actions/setup-example@v1 + + # Initializes the CodeQL tools for scanning. + - name: Initialize CodeQL + uses: github/codeql-action/init@v4 + with: + languages: ${{ matrix.language }} + build-mode: ${{ matrix.build-mode }} + # If you wish to specify custom queries, you can do so here or in a config file. + # By default, queries listed here will override any specified in a config file. + # Prefix the list here with "+" to use these queries and those in the config file. + + # For more details on CodeQL's query packs, refer to: https://docs.github.com/en/code-security/code-scanning/automatically-scanning-your-code-for-vulnerabilities-and-errors/configuring-code-scanning#using-queries-in-ql-packs + # queries: security-extended,security-and-quality + + # If the analyze step fails for one of the languages you are analyzing with + # "We were unable to automatically build your code", modify the matrix above + # to set the build mode to "manual" for that language. Then modify this step + # to build your code. + # ℹī¸ Command-line programs to run using the OS shell. + # 📚 See https://docs.github.com/en/actions/using-workflows/workflow-syntax-for-github-actions#jobsjob_idstepsrun + - name: Run manual build steps + if: matrix.build-mode == 'manual' + shell: bash + run: | + echo 'If you are using a "manual" build mode for one or more of the' \ + 'languages you are analyzing, replace this with the commands to build' \ + 'your code, for example:' + echo ' make bootstrap' + echo ' make release' + exit 1 + + - name: Perform CodeQL Analysis + uses: github/codeql-action/analyze@v4 + with: + category: "/language:${{matrix.language}}" diff --git a/project/block_manager/services/codegen/__init__.py b/project/block_manager/services/codegen/__init__.py new file mode 100644 index 0000000..5aad8b1 --- /dev/null +++ b/project/block_manager/services/codegen/__init__.py @@ -0,0 +1,6 @@ +"""Code generation orchestration package""" + +from .pytorch_orchestrator import PyTorchCodeOrchestrator +from .tensorflow_orchestrator import TensorFlowCodeOrchestrator + +__all__ = ['PyTorchCodeOrchestrator', 'TensorFlowCodeOrchestrator'] diff --git a/project/block_manager/services/codegen/base.py b/project/block_manager/services/codegen/base.py new file mode 100644 index 0000000..1c3d909 --- /dev/null +++ b/project/block_manager/services/codegen/base.py @@ -0,0 +1,79 @@ +""" +Base utilities for code generation +Shared functions for both PyTorch and TensorFlow code generation +""" + +from collections import deque +from typing import List, Dict, Any + + +def topological_sort(nodes: List[Dict], edges: List[Dict]) -> List[Dict]: + """ + Sort nodes in topological order based on edges using Kahn's algorithm. + + Args: + nodes: List of node definitions + edges: List of edge definitions + + Returns: + List of nodes in topological order + """ + node_map = {node['id']: node for node in nodes} + + # Build adjacency list and in-degree count + graph = {node['id']: [] for node in nodes} + in_degree = {node['id']: 0 for node in nodes} + + for edge in edges: + source = edge.get('source') + target = edge.get('target') + if source in graph and target in graph: + graph[source].append(target) + in_degree[target] += 1 + + # Kahn's algorithm + queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0]) + sorted_ids = [] + + while queue: + node_id = queue.popleft() + sorted_ids.append(node_id) + + for neighbor in graph[node_id]: + in_degree[neighbor] -= 1 + if in_degree[neighbor] == 0: + queue.append(neighbor) + + # Return nodes in sorted order + return [node_map[node_id] for node_id in sorted_ids if node_id in node_map] + + +def get_input_variable(incoming: List[str], var_map: Dict[str, str]) -> str: + """ + Determine input variable name based on incoming connections. + + Args: + incoming: List of incoming node IDs + var_map: Map of node ID to variable name + + Returns: + Variable name or list of variable names for multiple inputs + """ + if not incoming: + return 'x' + elif len(incoming) == 1: + return var_map.get(incoming[0], 'x') + else: + # Multiple inputs (for concat, add, etc.) + input_vars = [var_map.get(src, 'x') for src in incoming] + return f"[{', '.join(input_vars)}]" + + +def get_node_type(node: Dict[str, Any]) -> str: + """Extract node type from node definition""" + return node.get('data', {}).get('blockType', 'unknown') + + +def get_node_config(node: Dict[str, Any]) -> Dict[str, Any]: + """Extract configuration from node definition""" + return node.get('data', {}).get('config', {}) diff --git a/project/block_manager/services/codegen/base_orchestrator.py b/project/block_manager/services/codegen/base_orchestrator.py new file mode 100644 index 0000000..835c959 --- /dev/null +++ b/project/block_manager/services/codegen/base_orchestrator.py @@ -0,0 +1,339 @@ +""" +Base Code Generation Orchestrator +Shared functionality between PyTorch and TensorFlow orchestrators +""" + +from typing import List, Dict, Any, Optional, Tuple, Set +from collections import defaultdict +import json +from abc import ABC, abstractmethod + +from .base import topological_sort, get_input_variable, get_node_type, get_node_config +from ..nodes.registry import get_node_definition +from ..nodes.base import Framework, LayerCodeSpec +from ..nodes.templates.manager import TemplateManager + + +class UnsupportedNodeTypeError(Exception): + """Raised when a node type is not supported""" + pass + + +class BaseCodeOrchestrator(ABC): + """ + Base orchestrator for code generation. + Provides common functionality for both PyTorch and TensorFlow. + """ + + def __init__(self): + self.template_manager = TemplateManager() + + @property + @abstractmethod + def framework(self) -> Framework: + """Return the framework this orchestrator targets""" + pass + + @abstractmethod + def _get_code_spec_method_name(self) -> str: + """Return the method name for getting code specs (e.g., 'get_pytorch_code_spec')""" + pass + + @abstractmethod + def _generate_layer_call( + self, + layer_var: str, + input_var: str, + node_type: str, + spec: LayerCodeSpec + ) -> str: + """Generate the code for calling a layer (framework-specific)""" + pass + + @abstractmethod + def _get_default_input_shape(self) -> Tuple[int, ...]: + """Get default input shape for this framework""" + pass + + def generate( + self, + nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]], + project_name: str = "GeneratedModel", + group_definitions: Optional[List[Dict[str, Any]]] = None + ) -> Tuple[Dict[str, str], List[Exception]]: + """ + Generate complete project files. + + Args: + nodes: List of node definitions from the frontend + edges: List of edge definitions + project_name: Name for the generated model class + group_definitions: Optional group definitions + + Returns: + Tuple of (files dict, errors list) + """ + errors = [] + + try: + sorted_nodes = topological_sort(nodes, edges) + edge_map = self._build_edge_map(edges) + code_specs, spec_errors = self._generate_code_specs(sorted_nodes, edge_map) + errors.extend(spec_errors) + + layer_classes = self._render_layer_classes(code_specs) + model_definition = self._generate_model_definition( + project_name, code_specs, sorted_nodes, edge_map + ) + + input_shape = self._extract_input_shape(nodes) + test_code = self._generate_test_code(project_name, input_shape) + model_code = self._render_model_file( + project_name, layer_classes, model_definition, test_code + ) + + train_code = self._generate_training_script(project_name, nodes) + dataset_code = self._generate_dataset_script(nodes) + config_code = self._generate_config_file(nodes) + + return { + 'model': model_code, + 'train': train_code, + 'dataset': dataset_code, + 'config': config_code + }, errors + + except Exception as e: + errors.append(e) + return {}, errors + + def _build_edge_map(self, edges: List[Dict[str, Any]]) -> Dict[str, List[str]]: + """Build a map of node_id -> list of incoming node_ids""" + edge_map = defaultdict(list) + for edge in edges: + target = edge.get('target') + source = edge.get('source') + if target and source: + edge_map[target].append(source) + return dict(edge_map) + + def _generate_code_specs( + self, + sorted_nodes: List[Dict[str, Any]], + edge_map: Dict[str, List[str]] + ) -> Tuple[List[LayerCodeSpec], List[Exception]]: + """Generate code specifications for all nodes""" + code_specs = [] + errors = [] + + processable_nodes = [ + n for n in sorted_nodes + if get_node_type(n) not in ('input', 'dataloader', 'output') + ] + + for node in processable_nodes: + try: + node_id = node['id'] + node_type = get_node_type(node) + config = get_node_config(node) + + node_def = get_node_definition(node_type, self.framework) + + if not node_def: + raise UnsupportedNodeTypeError( + f"Node type '{node_type}' (id: {node_id}) is not supported for {self.framework.value}" + ) + + # Call the appropriate method dynamically + method = getattr(node_def, self._get_code_spec_method_name()) + code_spec = method( + node_id=node_id, + config=config, + input_shape=None, + output_shape=None + ) + + code_specs.append(code_spec) + + except Exception as e: + errors.append(e) + + return code_specs, errors + + def _render_layer_classes(self, code_specs: List[LayerCodeSpec]) -> str: + """Render all unique layer class definitions""" + unique_classes = {} + + for spec in code_specs: + if spec.node_type not in unique_classes: + try: + template_path = spec.get_template_path(self.framework) + rendered = self.template_manager.render( + template_path, + spec.template_context + ) + unique_classes[spec.node_type] = rendered + except Exception: + pass + + return '\n\n'.join(unique_classes.values()) + + def _generate_forward_pass( + self, + sorted_nodes: List[Dict[str, Any]], + edge_map: Dict[str, List[str]], + code_specs: List[LayerCodeSpec] + ) -> Tuple[List[str], Set[str]]: + """Generate forward pass logic with skip connection support""" + forward_lines = [] + var_map = {} + skip_connections = set() + spec_map = {spec.node_id: spec for spec in code_specs} + + processable_nodes = [ + n for n in sorted_nodes + if get_node_type(n) not in ('output',) + ] + + for node in processable_nodes: + node_id = node['id'] + node_type = get_node_type(node) + + if node_type in ('input', 'dataloader'): + var_map[node_id] = 'x' + continue + + incoming = edge_map.get(node_id, []) + input_var = get_input_variable(incoming, var_map) + + spec = spec_map.get(node_id) + if not spec: + continue + + output_var = f"x_{node_id.replace('-', '_')}" + + # Generate the layer call (framework-specific) + layer_call = self._generate_layer_call( + spec.layer_variable_name, + input_var, + node_type, + spec + ) + forward_lines.append(f"{output_var} = {layer_call}") + + var_map[node_id] = output_var + + if len(incoming) > 1: + skip_connections.add(output_var) + + # Ensure final output is assigned to 'x' + if processable_nodes: + last_node_id = processable_nodes[-1]['id'] + last_var = var_map.get(last_node_id, 'x') + if last_var != 'x': + forward_lines.append(f"x = {last_var}") + + return forward_lines, skip_connections + + def _extract_input_shape(self, nodes: List[Dict[str, Any]]) -> Tuple[int, ...]: + """Extract input shape from input node""" + input_node = next((n for n in nodes if get_node_type(n) == 'input'), None) + + if input_node: + config = get_node_config(input_node) + shape_str = config.get('shape', '') + try: + shape = json.loads(shape_str) if isinstance(shape_str, str) else shape_str + if isinstance(shape, list): + return tuple(shape) + except (ValueError, TypeError): + pass + + return self._get_default_input_shape() + + def _generate_training_script(self, project_name: str, nodes: List[Dict[str, Any]]) -> str: + """Generate training script using template""" + has_softmax = any(get_node_type(n) == 'softmax' for n in nodes) + is_classification = has_softmax + + context = self._get_training_context(project_name, is_classification) + template_path = f"{self.framework.value}/files/train.py.jinja2" + return self.template_manager.render(template_path, context) + + def _generate_dataset_script(self, nodes: List[Dict[str, Any]]) -> str: + """Generate dataset script using template""" + input_shape = self._extract_input_shape(nodes) + context = self._get_dataset_context(input_shape) + template_path = f"{self.framework.value}/files/dataset.py.jinja2" + return self.template_manager.render(template_path, context) + + def _generate_config_file(self, nodes: List[Dict[str, Any]]) -> str: + """Generate config file using template""" + input_shape = self._extract_input_shape(nodes) + layer_count = sum( + 1 for n in nodes + if get_node_type(n) not in ('input', 'output', 'dataloader') + ) + + if layer_count > 20: + batch_size, learning_rate, epochs, complexity = 16, 1e-4, 100, "Deep" + elif layer_count > 10: + batch_size, learning_rate, epochs, complexity = 32, 1e-3, 50, "Medium" + else: + batch_size, learning_rate, epochs, complexity = 64, 1e-3, 30, "Shallow" + + has_attention = any(get_node_type(n) in ('self_attention', 'attention') for n in nodes) + if has_attention: + learning_rate *= 0.1 + batch_size = max(8, batch_size // 2) + + context = { + 'batch_size': batch_size, + 'learning_rate': learning_rate, + 'num_epochs': epochs, + 'input_shape': list(input_shape), + 'complexity': complexity, + 'layer_count': layer_count, + 'has_attention': has_attention + } + + template_path = f"{self.framework.value}/files/config.py.jinja2" + return self.template_manager.render(template_path, context) + + @abstractmethod + def _generate_model_definition( + self, + project_name: str, + code_specs: List[LayerCodeSpec], + sorted_nodes: List[Dict[str, Any]], + edge_map: Dict[str, List[str]] + ) -> str: + """Generate the main model class definition (framework-specific)""" + pass + + @abstractmethod + def _generate_test_code(self, project_name: str, input_shape: Tuple[int, ...]) -> str: + """Generate test code (framework-specific)""" + pass + + @abstractmethod + def _render_model_file( + self, + project_name: str, + layer_classes: str, + model_definition: str, + test_code: str + ) -> str: + """Render the complete model file (framework-specific)""" + pass + + @abstractmethod + def _get_training_context(self, project_name: str, is_classification: bool) -> Dict[str, Any]: + """Get template context for training script (framework-specific)""" + pass + + @abstractmethod + def _get_dataset_context(self, input_shape: Tuple[int, ...]) -> Dict[str, Any]: + """Get template context for dataset script (framework-specific)""" + pass diff --git a/project/block_manager/services/codegen/group_block_generator.py b/project/block_manager/services/codegen/group_block_generator.py new file mode 100644 index 0000000..147ced0 --- /dev/null +++ b/project/block_manager/services/codegen/group_block_generator.py @@ -0,0 +1,192 @@ +""" +Group Block Code Generator +Abstract base class for generating group block code across frameworks. +""" + +from abc import ABC, abstractmethod +from typing import Dict, List, Any, Optional, Tuple +import re + +from ..nodes.base import Framework, LayerCodeSpec +from ..nodes.templates.manager import TemplateManager +from .base import topological_sort + + +class GroupBlockGenerator(ABC): + """ + Abstract base class for generating group block code across frameworks. + + Responsibilities: + - Parse group definition internal structure + - Detect multi-I/O patterns from portMappings + - Handle dependency ordering for nested internal nodes + - Coordinate with TemplateManager for rendering + """ + + def __init__(self, framework: Framework): + self.framework = framework + self.template_manager = TemplateManager() + + @abstractmethod + def generate_group_block_spec( + self, + group_definition: Dict[str, Any], + node_id: str, + instance_config: Optional[Dict[str, Any]] = None + ) -> LayerCodeSpec: + """ + Generate LayerCodeSpec for a group block instance. + + Args: + group_definition: The GroupBlockDefinition dict with internal_structure + node_id: The node ID of the group instance in the main graph + instance_config: Optional per-instance config overrides + + Returns: + LayerCodeSpec containing all info needed to render the group block + """ + pass + + @abstractmethod + def generate_group_class_code( + self, + group_definition: Dict[str, Any] + ) -> str: + """ + Generate the complete class code for a group block definition. + + Args: + group_definition: The GroupBlockDefinition dict + + Returns: + Rendered class code as string + """ + pass + + def _parse_port_mappings( + self, + port_mappings: List[Dict[str, Any]] + ) -> Tuple[List[Dict], List[Dict]]: + """ + Parse port mappings into input and output ports. + + Args: + port_mappings: List of port mapping dicts from internal_structure + + Returns: + Tuple of (input_ports, output_ports) + """ + input_ports = [pm for pm in port_mappings if pm.get('type') == 'input'] + output_ports = [pm for pm in port_mappings if pm.get('type') == 'output'] + return input_ports, output_ports + + def _topologically_sort_internal_nodes( + self, + nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """ + Topologically sort internal nodes using existing base.topological_sort. + + Args: + nodes: List of internal nodes + edges: List of internal edges + + Returns: + Topologically sorted list of nodes + """ + return topological_sort(nodes, edges) + + def _detect_nested_groups( + self, + internal_nodes: List[Dict[str, Any]] + ) -> List[str]: + """ + Detect which internal nodes are themselves group blocks. + + Args: + internal_nodes: List of internal nodes + + Returns: + List of group definition IDs referenced by internal nodes + """ + nested_groups = [] + for node in internal_nodes: + node_type = node.get('data', {}).get('blockType') + if node_type == 'group': + group_def_id = node.get('data', {}).get('groupDefinitionId') + if group_def_id: + nested_groups.append(group_def_id) + return nested_groups + + def _build_template_context( + self, + group_definition: Dict[str, Any], + internal_nodes: List[Dict[str, Any]], + input_ports: List[Dict], + output_ports: List[Dict] + ) -> Dict[str, Any]: + """ + Build the base template context for rendering group block class. + + Args: + group_definition: The group definition dict + internal_nodes: Sorted list of internal nodes + input_ports: List of input port mappings + output_ports: List of output port mappings + + Returns: + Dict with keys: class_name, description, layers, input_ports, + output_ports, has_multi_input, has_multi_output, etc. + """ + return { + 'class_name': self._sanitize_class_name(group_definition['name']), + 'group_name': group_definition['name'], + 'description': group_definition.get('description', ''), + 'input_ports': input_ports, + 'output_ports': output_ports, + 'has_multi_input': len(input_ports) > 1, + 'has_multi_output': len(output_ports) > 1, + 'num_inputs': len(input_ports), + 'num_outputs': len(output_ports), + 'internal_nodes': internal_nodes, + 'framework': self.framework.value + } + + def _sanitize_class_name(self, name: str) -> str: + """ + Convert group name to valid class name (PascalCase). + + Args: + name: Original group name (may contain spaces, special chars) + + Returns: + Valid PascalCase class name + """ + # Remove special chars, capitalize words + clean = re.sub(r'[^a-zA-Z0-9_]', '', name.replace(' ', '_')) + parts = clean.split('_') + return ''.join(word.capitalize() for word in parts if word) + + def _build_edge_map( + self, + edges: List[Dict[str, Any]] + ) -> Dict[str, List[str]]: + """ + Build a map of node_id -> list of incoming node_ids. + + Args: + edges: List of edge dicts + + Returns: + Dict mapping target node ID to list of source node IDs + """ + edge_map = {} + for edge in edges: + target = edge.get('target') + source = edge.get('source') + if target and source: + if target not in edge_map: + edge_map[target] = [] + edge_map[target].append(source) + return edge_map diff --git a/project/block_manager/services/codegen/pytorch_group_generator.py b/project/block_manager/services/codegen/pytorch_group_generator.py new file mode 100644 index 0000000..a03f560 --- /dev/null +++ b/project/block_manager/services/codegen/pytorch_group_generator.py @@ -0,0 +1,397 @@ +""" +PyTorch Group Block Generator +Generates PyTorch nn.Module code for group block definitions. +""" + +from typing import Dict, List, Any, Optional +from .group_block_generator import GroupBlockGenerator +from ..nodes.base import Framework, LayerCodeSpec +from ..nodes.registry import get_node_definition +from .base import get_node_type, get_node_config + + +class PyTorchGroupBlockGenerator(GroupBlockGenerator): + """PyTorch-specific group block code generation""" + + def __init__(self): + super().__init__(Framework.PYTORCH) + + def generate_group_block_spec( + self, + group_definition: Dict[str, Any], + node_id: str, + instance_config: Optional[Dict[str, Any]] = None, + input_shape: Optional[Any] = None + ) -> LayerCodeSpec: + """ + Generate LayerCodeSpec for PyTorch group block instance. + + Args: + group_definition: The GroupBlockDefinition dict + node_id: The node ID of the group instance in main graph + instance_config: Optional per-instance config overrides + input_shape: Optional input shape for the group block + + Returns: + LayerCodeSpec for this group instance + """ + class_name = self._sanitize_class_name(group_definition['name']) + sanitized_id = node_id.replace('-', '_') + layer_var_name = f"{sanitized_id}_{class_name}" + + # Extract init params from internal structure with shape information + init_params = self._extract_init_params(group_definition, instance_config, input_shape) + + # Build template context + internal_structure = group_definition.get('internal_structure', {}) + port_mappings = internal_structure.get('portMappings', []) + input_ports, output_ports = self._parse_port_mappings(port_mappings) + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var_name, + node_type='group', + node_id=node_id, + init_params=init_params, + config_params=instance_config or {}, + template_context={ + 'group_definition_id': group_definition['id'], + 'has_multi_output': len(output_ports) > 1, + 'num_outputs': len(output_ports), + 'num_inputs': len(input_ports) + } + ) + + def generate_group_class_code( + self, + group_definition: Dict[str, Any], + input_shape: Optional[Any] = None + ) -> str: + """ + Generate PyTorch group block class using template. + + Args: + group_definition: The GroupBlockDefinition dict + input_shape: Optional representative input shape for the group block + + Returns: + Rendered class code as string + """ + internal_structure = group_definition.get('internal_structure', {}) + internal_nodes = internal_structure.get('nodes', []) + internal_edges = internal_structure.get('edges', []) + port_mappings = internal_structure.get('portMappings', []) + + # Sort internal nodes + sorted_nodes = self._topologically_sort_internal_nodes( + internal_nodes, internal_edges + ) + + # Parse ports + input_ports, output_ports = self._parse_port_mappings(port_mappings) + + # Generate LayerCodeSpecs for internal nodes with shape inference + internal_specs = self._generate_internal_node_specs( + sorted_nodes, + internal_edges=internal_edges, + input_shape=input_shape + ) + + # Build edge map + edge_map = self._build_edge_map(internal_edges) + + # Generate forward pass lines + forward_lines = self._generate_forward_pass( + sorted_nodes, + internal_specs, + edge_map, + input_ports, + output_ports + ) + + # Build template context + context = self._build_pytorch_template_context( + group_definition, + sorted_nodes, + internal_specs, + input_ports, + output_ports, + forward_lines + ) + + # Render template + template_path = 'pytorch/layers/group_block.py.jinja2' + return self.template_manager.render(template_path, context) + + def _generate_internal_node_specs( + self, + internal_nodes: List[Dict[str, Any]], + internal_edges: Optional[List[Dict[str, Any]]] = None, + input_shape: Optional[Any] = None + ) -> List[LayerCodeSpec]: + """ + Generate code specs for each internal node with shape inference. + + Args: + internal_nodes: Sorted list of internal nodes + internal_edges: List of internal edges for shape propagation + input_shape: Input shape to the group block + + Returns: + List of LayerCodeSpec for processable internal nodes + """ + from ..nodes.rules.shape import TensorShape + + specs = [] + + # Build edge map for shape inference + edge_map = {} + if internal_edges: + from collections import defaultdict + edge_map_builder = defaultdict(list) + for edge in internal_edges: + target = edge.get('target') + source = edge.get('source') + if target and source: + edge_map_builder[target].append(source) + edge_map = dict(edge_map_builder) + + # Track output shapes of each node + node_output_shapes = {} + + # Initialize input nodes with the group's input shape + for node in internal_nodes: + node_type = get_node_type(node) + if node_type == 'input': + node_output_shapes[node['id']] = input_shape + + for node in internal_nodes: + node_type = get_node_type(node) + + # Skip special nodes + if node_type in ('input', 'output', 'dataloader'): + continue + + node_id = node['id'] + config = get_node_config(node) + + # Check if this is a nested group + if node_type == 'group': + # For Phase 1-2, we can defer nested groups + # TODO: Implement nested group support + continue + + # Determine input shape from incoming connections + computed_input_shape = None + incoming = edge_map.get(node_id, []) + if incoming: + if len(incoming) == 1: + computed_input_shape = node_output_shapes.get(incoming[0]) + else: + # Multiple inputs - use first for now + for src in incoming: + if src in node_output_shapes: + computed_input_shape = node_output_shapes[src] + break + + # Get node definition from registry + node_def = get_node_definition(node_type, Framework.PYTORCH) + if node_def: + # Compute output shape + computed_output_shape = None + try: + if hasattr(node_def, 'compute_output_shape'): + computed_output_shape = node_def.compute_output_shape(computed_input_shape, config) + except Exception: + pass + + # Generate spec with shape information + spec = node_def.get_pytorch_code_spec( + node_id=node_id, + config=config, + input_shape=computed_input_shape, + output_shape=computed_output_shape + ) + specs.append(spec) + + # Store output shape for downstream nodes + if computed_output_shape: + node_output_shapes[node_id] = computed_output_shape + + return specs + + def _generate_forward_pass( + self, + sorted_nodes: List[Dict[str, Any]], + internal_specs: List[LayerCodeSpec], + edge_map: Dict[str, List[str]], + input_ports: List[Dict], + output_ports: List[Dict] + ) -> List[str]: + """ + Generate forward pass lines for internal graph. + + Args: + sorted_nodes: Topologically sorted internal nodes + internal_specs: LayerCodeSpecs for internal nodes + edge_map: Map of node_id -> [incoming_node_ids] + input_ports: Input port mappings + output_ports: Output port mappings + + Returns: + List of forward pass code lines + """ + forward_lines = [] + var_map = {} + spec_map = {spec.node_id: spec for spec in internal_specs} + + # Map input ports to their external parameter names + input_node_map = { + port['internalNodeId']: port + for port in input_ports + } + + # Track which nodes are outputs (need to preserve their variables) + output_node_ids = {port['internalNodeId'] for port in output_ports} + + for node in sorted_nodes: + node_id = node['id'] + node_type = get_node_type(node) + + # Handle input nodes + if node_type == 'input': + # Map input node to its external parameter + if node_id in input_node_map: + port_info = input_node_map[node_id] + # For now, use simple variable naming + if len(input_ports) == 1: + var_map[node_id] = 'x' + else: + # Multi-input: use port labels or indices + port_label = port_info.get('externalPortLabel', f'input_{len(var_map)}') + var_name = port_label.lower().replace(' ', '_') + var_map[node_id] = var_name + continue + + # Skip output and dataloader nodes (they don't produce code) + if node_type in ('output', 'dataloader'): + continue + + # Get the spec for this node + spec = spec_map.get(node_id) + if not spec: + # Node not in specs (might be nested group or unsupported) + continue + + # Determine input variable + incoming = edge_map.get(node_id, []) + if not incoming: + input_var = 'x' + elif len(incoming) == 1: + input_var = var_map.get(incoming[0], 'x') + else: + # Multiple inputs (add, concat nodes) + input_vars = [var_map.get(src, 'x') for src in incoming] + input_var = f"[{', '.join(input_vars)}]" + + # Generate output variable + output_var = f"x_{node_id.replace('-', '_')}" + + # Generate layer call + if node_type in ('add', 'concat'): + # Special handling for merge nodes + if node_type == 'concat': + dim = spec.template_context.get('dim', 1) + forward_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var}, concat_dim={dim})" + ) + else: + forward_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var})" + ) + else: + # Standard layer call + forward_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var})" + ) + + var_map[node_id] = output_var + + return forward_lines + + def _build_pytorch_template_context( + self, + group_definition: Dict[str, Any], + sorted_nodes: List[Dict[str, Any]], + internal_specs: List[LayerCodeSpec], + input_ports: List[Dict], + output_ports: List[Dict], + forward_lines: List[str] + ) -> Dict[str, Any]: + """ + Build comprehensive template context for PyTorch group block. + + Args: + group_definition: The group definition + sorted_nodes: Sorted internal nodes + internal_specs: LayerCodeSpecs for internal nodes + input_ports: Input port mappings + output_ports: Output port mappings + forward_lines: Generated forward pass code lines + + Returns: + Complete template context dict + """ + base_context = self._build_template_context( + group_definition, sorted_nodes, input_ports, output_ports + ) + + # Build input parameter signature + if len(input_ports) == 1: + input_params = 'x' + else: + # Multi-input: generate parameter names from port labels + param_names = [] + for port in input_ports: + label = port.get('externalPortLabel', f'input_{len(param_names)}') + param_name = label.lower().replace(' ', '_') + param_names.append(param_name) + input_params = ', '.join(param_names) + + # Build output variables + output_node_ids = [port['internalNodeId'] for port in output_ports] + output_vars = [] + for node_id in output_node_ids: + # Use the variable name from forward pass + var_name = f"x_{node_id.replace('-', '_')}" + output_vars.append(var_name) + + return { + **base_context, + 'internal_specs': internal_specs, + 'forward_lines': forward_lines, + 'input_params': input_params, + 'output_vars': output_vars, + 'init_params': [] # No custom init params for Phase 1-2 + } + + def _extract_init_params( + self, + group_definition: Dict[str, Any], + instance_config: Optional[Dict[str, Any]], + input_shape: Optional[Any] = None + ) -> Dict[str, Any]: + """ + Extract initialization parameters for group block instance. + + Args: + group_definition: The group definition + instance_config: Per-instance config overrides + input_shape: Optional input shape for the group block + + Returns: + Dict of init parameters (empty for Phase 1-2) + """ + # For Phase 1-2, group blocks don't have instance-level init params + # All config is baked into the class definition + return {} diff --git a/project/block_manager/services/codegen/pytorch_orchestrator.py b/project/block_manager/services/codegen/pytorch_orchestrator.py new file mode 100644 index 0000000..39aef55 --- /dev/null +++ b/project/block_manager/services/codegen/pytorch_orchestrator.py @@ -0,0 +1,928 @@ +""" +PyTorch Code Generation Orchestrator +Coordinates the generation of complete PyTorch project files +""" + +from typing import List, Dict, Any, Optional, Tuple, Set +from collections import defaultdict +import json + +from .base import topological_sort, get_input_variable, get_node_type, get_node_config +from ..nodes.registry import get_node_definition +from ..nodes.base import Framework, LayerCodeSpec +from ..nodes.templates.manager import TemplateManager + + +class UnsupportedNodeTypeError(Exception): + """Raised when a node type is not supported""" + pass + + +class PyTorchCodeOrchestrator: + """ + Orchestrator for PyTorch code generation. + Delegates code generation to individual node classes and assembles the final output. + """ + + def __init__(self): + self.template_manager = TemplateManager() + + def generate( + self, + nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]], + project_name: str = "GeneratedModel", + group_definitions: Optional[List[Dict[str, Any]]] = None + ) -> Tuple[Dict[str, str], List[Exception]]: + """ + Generate complete PyTorch project files. + + Args: + nodes: List of node definitions from the frontend + edges: List of edge definitions + project_name: Name for the generated model class + group_definitions: Optional list of group block definitions + + Returns: + Tuple of (files dict, errors list) + files dict contains: {'model': str, 'train': str, 'dataset': str, 'config': str} + """ + errors = [] + + try: + # Initialize group block generator if needed + group_generator = None + if group_definitions: + from .pytorch_group_generator import PyTorchGroupBlockGenerator + group_generator = PyTorchGroupBlockGenerator() + + # Sort nodes topologically + sorted_nodes = topological_sort(nodes, edges) + + # Build edge map for quick lookups + edge_map = self._build_edge_map(edges) + + # Generate code specifications for each node + code_specs, spec_errors = self._generate_code_specs( + sorted_nodes, edge_map, group_generator, group_definitions + ) + errors.extend(spec_errors) + + # Generate code specs for internal layers in group blocks + if group_definitions: + internal_specs, internal_errors = self._generate_internal_layer_specs( + group_definitions + ) + code_specs.extend(internal_specs) + errors.extend(internal_errors) + + # Render layer classes from templates (includes internal layers) + layer_classes = self._render_layer_classes(code_specs) + + # Generate group block class definitions + group_classes = "" + if group_generator and group_definitions: + group_classes = self._generate_group_block_classes( + group_definitions, group_generator, sorted_nodes, edge_map + ) + + # Combine regular layers + group classes + all_classes = layer_classes + if group_classes: + all_classes += "\n\n" + group_classes + + # Generate model class definition + model_definition = self._generate_model_definition( + project_name, + code_specs, + sorted_nodes, + edges, + edge_map + ) + + # Generate test code + input_shape = self._extract_input_shape(nodes) + test_code = self._generate_test_code(project_name, input_shape) + + # Render complete model file + model_code = self._render_model_file( + project_name, + all_classes, + model_definition, + test_code + ) + + # Generate training script + train_code = self._generate_training_script(project_name, nodes) + + # Generate dataset script + dataset_code = self._generate_dataset_script(nodes) + + # Generate config file + config_code = self._generate_config_file(nodes) + + return { + 'model': model_code, + 'train': train_code, + 'dataset': dataset_code, + 'config': config_code + }, errors + + except Exception as e: + errors.append(e) + return {}, errors + + def _build_edge_map(self, edges: List[Dict[str, Any]]) -> Dict[str, List[str]]: + """Build a map of node_id -> list of incoming node_ids""" + edge_map = defaultdict(list) + for edge in edges: + target = edge.get('target') + source = edge.get('source') + if target and source: + edge_map[target].append(source) + return dict(edge_map) + + def _build_outgoing_edge_map(self, edges: List[Dict[str, Any]]) -> Dict[str, List[str]]: + """ + Build a map of outgoing edges from each node. + + Returns: + Dict mapping source_node_id -> [target_node_ids] + """ + outgoing_map = defaultdict(list) + for edge in edges: + source = edge.get('source') + target = edge.get('target') + if source and target: + outgoing_map[source].append(target) + return dict(outgoing_map) + + def _compute_shape_map( + self, + sorted_nodes: List[Dict[str, Any]], + edge_map: Dict[str, List[str]], + group_definitions: Optional[List[Dict[str, Any]]] = None + ) -> Dict[str, Any]: + """ + Compute input and output shapes for all nodes through forward propagation. + + Returns: + Dict mapping "{node_id}_input" and "{node_id}_output" to TensorShape + """ + from ..nodes.rules.shape import TensorShape + + shape_map = {} + node_output_shapes = {} # Track each node's output shape + + for node in sorted_nodes: + node_id = node['id'] + node_type = get_node_type(node) + config = get_node_config(node) + + # Handle input nodes + if node_type == 'input': + # Extract shape from input node config + shape_str = config.get('shape', '[1, 3, 224, 224]') + try: + import json + shape_list = json.loads(shape_str) if isinstance(shape_str, str) else shape_str + if isinstance(shape_list, list): + output_shape = TensorShape({'dims': shape_list, 'description': 'Input'}) + node_output_shapes[node_id] = output_shape + except (ValueError, TypeError): + # Default shape if parsing fails + node_output_shapes[node_id] = TensorShape({'dims': [1, 3, 224, 224], 'description': 'Input'}) + continue + + # Handle dataloader nodes + if node_type == 'dataloader': + # Dataloader typically outputs batched image data + # Use default or extract from config if available + shape_str = config.get('output_shape', '[1, 3, 224, 224]') + try: + import json + shape_list = json.loads(shape_str) if isinstance(shape_str, str) else shape_str + if isinstance(shape_list, list): + output_shape = TensorShape({'dims': shape_list, 'description': 'Dataloader output'}) + node_output_shapes[node_id] = output_shape + except (ValueError, TypeError): + node_output_shapes[node_id] = TensorShape({'dims': [1, 3, 224, 224], 'description': 'Dataloader output'}) + continue + + # Skip output nodes + if node_type == 'output': + continue + + # Get incoming nodes + incoming = edge_map.get(node_id, []) + + # Determine input shape from incoming connections + input_shape = None + if incoming: + if len(incoming) == 1: + # Single input + input_shape = node_output_shapes.get(incoming[0]) + else: + # Multiple inputs (for add, concat nodes) + # Store all input shapes for multi-input operations + input_shapes = [node_output_shapes.get(src) for src in incoming] + input_shapes = [s for s in input_shapes if s is not None] + if input_shapes: + # For now, use the first input shape as primary + # Specific nodes (add, concat) will handle multiple shapes + input_shape = input_shapes[0] + + # Store input shape in map + if input_shape: + shape_map[f"{node_id}_input"] = input_shape + + # Compute output shape using node definition + output_shape = None + try: + if node_type == 'group': + # For group blocks, we'll need to infer from internal structure + # For now, pass through input shape (placeholder) + output_shape = input_shape + else: + node_def = get_node_definition(node_type, Framework.PYTORCH) + if node_def and hasattr(node_def, 'compute_output_shape'): + output_shape = node_def.compute_output_shape(input_shape, config) + except Exception: + # If shape computation fails, use None + pass + + # Store output shape + if output_shape: + shape_map[f"{node_id}_output"] = output_shape + node_output_shapes[node_id] = output_shape + + return shape_map + + def _generate_code_specs( + self, + sorted_nodes: List[Dict[str, Any]], + edge_map: Dict[str, List[str]], + group_generator=None, + group_definitions: Optional[List[Dict[str, Any]]] = None + ) -> Tuple[List[LayerCodeSpec], List[Exception]]: + """ + Generate code specifications for all nodes including group blocks. + + Returns: + Tuple of (list of code specs, list of errors) + """ + code_specs = [] + errors = [] + + # Build group definition lookup + group_def_map = {} + if group_definitions: + group_def_map = {gd['id']: gd for gd in group_definitions} + + # Compute shape map for all nodes + shape_map = self._compute_shape_map(sorted_nodes, edge_map, group_definitions) + + # Skip input/dataloader/output nodes - they don't generate layers + processable_nodes = [ + n for n in sorted_nodes + if get_node_type(n) not in ('input', 'dataloader', 'output') + ] + + for node in processable_nodes: + try: + node_id = node['id'] + node_type = get_node_type(node) + config = get_node_config(node) + + # Get input and output shapes from shape map + input_shape = shape_map.get(f"{node_id}_input") + output_shape = shape_map.get(f"{node_id}_output") + + # Handle group blocks + if node_type == 'group': + if not group_generator: + raise UnsupportedNodeTypeError( + f"Group node {node_id} found but no group_definitions provided" + ) + + group_def_id = node.get('data', {}).get('groupDefinitionId') + group_def = group_def_map.get(group_def_id) + + if not group_def: + raise ValueError( + f"Group definition {group_def_id} not found for node {node_id}" + ) + + # Generate spec for this group instance + code_spec = group_generator.generate_group_block_spec( + group_definition=group_def, + node_id=node_id, + instance_config=config, + input_shape=input_shape + ) + code_specs.append(code_spec) + + else: + # Regular node - existing logic + node_def = get_node_definition(node_type, Framework.PYTORCH) + + if not node_def: + raise UnsupportedNodeTypeError( + f"Node type '{node_type}' (id: {node_id}) is not supported for PyTorch" + ) + + # Generate code specification with shape information + code_spec = node_def.get_pytorch_code_spec( + node_id=node_id, + config=config, + input_shape=input_shape, + output_shape=output_shape + ) + + code_specs.append(code_spec) + + except Exception as e: + errors.append(e) + + return code_specs, errors + + def _generate_internal_layer_specs( + self, + group_definitions: List[Dict[str, Any]] + ) -> Tuple[List[LayerCodeSpec], List[Exception]]: + """ + Generate LayerCodeSpecs for all unique internal layers used in group blocks. + This ensures internal layer classes are defined before group blocks use them. + + Args: + group_definitions: List of group block definitions + + Returns: + Tuple of (list of internal layer specs, list of errors) + """ + internal_specs = [] + errors = [] + seen_node_types = set() + + for group_def in group_definitions: + internal_structure = group_def.get('internal_structure', {}) + internal_nodes = internal_structure.get('nodes', []) + + for node in internal_nodes: + node_type = get_node_type(node) + + # Skip special nodes + if node_type in ('input', 'output', 'dataloader', 'group'): + continue + + # Only generate each node type once + if node_type in seen_node_types: + continue + + seen_node_types.add(node_type) + + try: + node_id = node['id'] + config = get_node_config(node) + + # Get node definition from registry + node_def = get_node_definition(node_type, Framework.PYTORCH) + + if not node_def: + raise UnsupportedNodeTypeError( + f"Internal node type '{node_type}' not supported in group block" + ) + + # Generate code specification + code_spec = node_def.get_pytorch_code_spec( + node_id=node_id, + config=config, + input_shape=None, + output_shape=None + ) + + internal_specs.append(code_spec) + + except Exception as e: + errors.append(e) + + return internal_specs, errors + + def _render_layer_classes(self, code_specs: List[LayerCodeSpec]) -> str: + """ + Render all unique layer class definitions. + + Returns: + String containing all layer class definitions + """ + unique_classes = {} + + for spec in code_specs: + # Use node_type as key to ensure we only define each class type once + if spec.node_type not in unique_classes: + try: + template_path = spec.get_template_path(Framework.PYTORCH) + # Merge class_name, init_params, and template_context for rendering + context = { + 'class_name': spec.class_name, + **spec.init_params, + **spec.template_context + } + rendered = self.template_manager.render(template_path, context) + unique_classes[spec.node_type] = rendered + except Exception: + # If template doesn't exist, skip this layer + # Error will be caught during model generation + pass + + # Join all classes with blank lines + return '\n\n'.join(unique_classes.values()) + + def _generate_group_block_classes( + self, + group_definitions: List[Dict[str, Any]], + group_generator, + sorted_nodes: List[Dict[str, Any]], + edge_map: Dict[str, List[str]] + ) -> str: + """ + Generate class definitions for all group blocks with representative input shapes. + + Args: + group_definitions: List of group block definitions + group_generator: PyTorchGroupBlockGenerator instance + sorted_nodes: Sorted nodes from main graph + edge_map: Edge map for shape inference + + Returns: + String containing all group block class definitions + """ + # Compute shape map for the main graph + shape_map = self._compute_shape_map(sorted_nodes, edge_map, group_definitions) + + # Extract representative input shapes for each group definition + group_shape_map = self._extract_group_representative_shapes( + group_definitions, sorted_nodes, shape_map + ) + + # Detect dependency order for nested groups + ordered_definitions = self._order_group_definitions(group_definitions) + + class_codes = [] + for group_def in ordered_definitions: + try: + # Get representative input shape for this group definition + representative_shape = group_shape_map.get(group_def['id']) + + class_code = group_generator.generate_group_class_code( + group_def, + input_shape=representative_shape + ) + class_codes.append(class_code) + except Exception as e: + # Log error but continue with other groups + print(f"Error generating group {group_def.get('name')}: {e}") + + return '\n\n'.join(class_codes) + + def _extract_group_representative_shapes( + self, + group_definitions: List[Dict[str, Any]], + sorted_nodes: List[Dict[str, Any]], + shape_map: Dict[str, Any] + ) -> Dict[str, Any]: + """ + Extract representative input shapes for each group definition. + Finds the first usage of each group definition and uses its input shape. + + Args: + group_definitions: List of group block definitions + sorted_nodes: Sorted nodes from main graph + shape_map: Shape map containing input shapes for all nodes + + Returns: + Dict mapping group_definition_id -> representative input shape + """ + group_shape_map = {} + + for group_def in group_definitions: + group_def_id = group_def['id'] + + # Find first usage of this group definition in the main graph + for node in sorted_nodes: + if get_node_type(node) == 'group': + node_group_def_id = node.get('data', {}).get('groupDefinitionId') + if node_group_def_id == group_def_id: + # Found a usage - get its input shape + node_id = node['id'] + input_shape = shape_map.get(f"{node_id}_input") + if input_shape: + group_shape_map[group_def_id] = input_shape + break + + return group_shape_map + + def _order_group_definitions( + self, + group_definitions: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """ + Order group definitions so nested groups are defined before their parents. + Uses topological sort based on group dependencies. + + Args: + group_definitions: List of group block definitions + + Returns: + Topologically sorted list of group definitions + """ + # Build dependency graph + graph = {gd['id']: [] for gd in group_definitions} + + for group_def in group_definitions: + internal_nodes = group_def.get('internal_structure', {}).get('nodes', []) + for node in internal_nodes: + if node.get('data', {}).get('blockType') == 'group': + nested_group_id = node.get('data', {}).get('groupDefinitionId') + if nested_group_id in graph: + # This group depends on the nested group + graph[group_def['id']].append(nested_group_id) + + # Topological sort (simple implementation) + visited = set() + result = [] + + def visit(gd_id): + if gd_id in visited: + return + visited.add(gd_id) + for dep in graph.get(gd_id, []): + visit(dep) + result.append(gd_id) + + for gd in group_definitions: + visit(gd['id']) + + # Map back to group definitions + gd_map = {gd['id']: gd for gd in group_definitions} + return [gd_map[gd_id] for gd_id in result if gd_id in gd_map] + + def _generate_model_definition( + self, + project_name: str, + code_specs: List[LayerCodeSpec], + sorted_nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]], + edge_map: Dict[str, List[str]] + ) -> str: + """Generate the main model class definition""" + # Generate layer initializations + layer_inits = [] + for spec in code_specs: + params_str = ', '.join( + f"{k}={repr(v)}" for k, v in spec.init_params.items() + ) + layer_inits.append( + f"self.{spec.layer_variable_name} = {spec.class_name}({params_str})" + ) + + # Generate forward pass logic with skip connection support + forward_lines, skip_connections = self._generate_forward_pass( + sorted_nodes, + edges, + edge_map, + code_specs + ) + + model_class = f'''class {project_name}(nn.Module): + """ + PyTorch Model for {project_name} + + This model is auto-generated from the VisionForge architecture. + """ + + def __init__(self): + super({project_name}, self).__init__() + #========================== + #Layer Initializations: + #========================== +{chr(10).join(" " + line for line in layer_inits)} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the model. + + Args: + x: Input tensor + + Returns: + Output tensor after processing through the model + """ +{chr(10).join(" " + line for line in forward_lines)} + return x +''' + return model_class + + def _needs_named_variable( + self, + node_id: str, + node_type: str, + outgoing_edge_map: Dict[str, List[str]], + num_outputs: int + ) -> bool: + """ + Determine if a node's output requires a named variable. + + A named variable is needed when: + 1. Node has multiple outgoing edges (output used multiple times - skip connections) + 2. Node has multiple outputs (e.g., group blocks) + 3. Node is a merge operation (add, concat) - for readability + + Args: + node_id: Node identifier + node_type: Type of node (e.g., 'conv2d', 'add') + outgoing_edge_map: Map of node_id -> list of target node IDs + num_outputs: Number of outputs for this node + + Returns: + True if a named variable should be created + """ + # Multi-output nodes always need named variables + if num_outputs > 1: + return True + + # Nodes with multiple outgoing edges need named variables (skip connections) + outgoing_edges = outgoing_edge_map.get(node_id, []) + if len(outgoing_edges) > 1: + return True + + # Merge operations benefit from named variables for readability + if node_type in ('add', 'concat'): + return True + + return False + + def _generate_forward_pass( + self, + sorted_nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]], + edge_map: Dict[str, List[str]], + code_specs: List[LayerCodeSpec] + ) -> Tuple[List[str], Set[str]]: + """ + Generate forward pass logic with optimized variable usage. + + Only creates named variables when needed: + - Skip connections (nodes with multiple outgoing edges) + - Multi-output nodes (group blocks) + - Merge operations (add, concat) + + Otherwise reuses 'x' variable for memory efficiency. + + Returns: + Tuple of (forward pass lines, set of skip connection var names) + """ + forward_lines = [] + var_map = {} # Maps node_id to variable name + skip_connections = set() + spec_map = {spec.node_id: spec for spec in code_specs} + + # Build outgoing edge map to detect skip connections + outgoing_edge_map = self._build_outgoing_edge_map(edges) + + # Process nodes in topological order + processable_nodes = [ + n for n in sorted_nodes + if get_node_type(n) not in ('output',) # Keep input/dataloader for var mapping + ] + + for node in processable_nodes: + node_id = node['id'] + node_type = get_node_type(node) + + # Input and dataloader nodes set up the initial 'x' variable + if node_type in ('input', 'dataloader'): + var_map[node_id] = 'x' + continue + + # Get incoming connections + incoming = edge_map.get(node_id, []) + + # Determine input variable + input_var = get_input_variable(incoming, var_map) + + # Get code spec for this node + spec = spec_map.get(node_id) + if not spec: + continue + + # Handle group blocks with multiple outputs + if node_type == 'group' and spec.template_context.get('has_multi_output'): + num_outputs = spec.template_context.get('num_outputs', 1) + output_vars = [ + f"{node_id.replace('-', '_')}_out{i}" + for i in range(num_outputs) + ] + forward_lines.append( + f"{', '.join(output_vars)} = self.{spec.layer_variable_name}({input_var})" + ) + # Map first output as primary variable for this node + var_map[node_id] = output_vars[0] + skip_connections.add(output_vars[0]) + + else: + # Determine if this node needs a named variable + num_outputs = 1 + needs_named_var = self._needs_named_variable( + node_id, + node_type, + outgoing_edge_map, + num_outputs + ) + + if needs_named_var: + # Create named variable for skip connections, merge ops, etc. + output_var = f"x_{node_id.replace('-', '_')}" + + # Handle concat with dimension parameter + if node_type == 'concat': + dim = spec.template_context.get('dim', 1) + forward_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var}, concat_dim={dim})" + ) + else: + forward_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var})" + ) + + var_map[node_id] = output_var + + # Track if this is a skip connection source + if len(outgoing_edge_map.get(node_id, [])) > 1: + skip_connections.add(output_var) + + else: + # Reuse 'x' variable for linear chains (memory efficient) + forward_lines.append( + f"x = self.{spec.layer_variable_name}({input_var})" + ) + var_map[node_id] = 'x' + + # Ensure final output is assigned to 'x' for return statement + if processable_nodes: + last_node_id = processable_nodes[-1]['id'] + last_var = var_map.get(last_node_id, 'x') + if last_var != 'x': + forward_lines.append(f"x = {last_var}") + + return forward_lines, skip_connections + + def _extract_input_shape(self, nodes: List[Dict[str, Any]]) -> Tuple[int, ...]: + """Extract input shape from input node""" + input_node = next((n for n in nodes if get_node_type(n) == 'input'), None) + + if input_node: + config = get_node_config(input_node) + shape_str = config.get('shape', '[1, 3, 224, 224]') + try: + shape = json.loads(shape_str) if isinstance(shape_str, str) else shape_str + if isinstance(shape, list): + return tuple(shape) + except (ValueError, TypeError): + pass + + return (1, 3, 224, 224) + + def _generate_test_code(self, project_name: str, input_shape: Tuple[int, ...]) -> str: + """Generate test code for model validation""" + return f'''if __name__ == "__main__": + # Test the model with random input + model = {project_name}() + model.eval() + test_input = torch.randn({input_shape}) + print(f"Input shape: {{test_input.shape}}") + output = model(test_input) + print(f"Output shape: {{output.shape}}") + print(f"Model has {{sum(p.numel() for p in model.parameters()):,}} parameters") +''' + + def _render_model_file( + self, + project_name: str, + layer_classes: str, + model_definition: str, + test_code: str + ) -> str: + """Render the complete model.py file""" + context = { + 'project_name': project_name, + 'layer_classes': layer_classes, + 'model_class_name': project_name, + 'layer_initializations': [], # Handled in model_definition + 'forward_pass_lines': [], # Handled in model_definition + 'test_code': test_code + } + + # For now, use string formatting since we're embedding pre-rendered content + # TODO: Convert model_definition to use template as well + return f'''""" +Generated PyTorch Model +Architecture: {project_name} +Generated by VisionForge + +This file contains the model architecture with separate layer classes. +Each layer is implemented as a reusable class for clarity and maintainability. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List, Tuple, Optional + + +#========================== +#Layer Definitions: +#========================== +{layer_classes} + +{model_definition} + +{test_code} +''' + + def _generate_training_script(self, project_name: str, nodes: List[Dict[str, Any]]) -> str: + """Generate training script using template""" + # Determine task type based on architecture + has_softmax = any(get_node_type(n) == 'softmax' for n in nodes) + is_classification = has_softmax + + context = { + 'project_name': project_name, + 'model_class_name': project_name, + 'task_type': 'classification' if is_classification else 'regression', + 'is_classification': is_classification, + 'loss_function': 'nn.CrossEntropyLoss()' if is_classification else 'nn.MSELoss()', + 'metric_name': 'accuracy' if is_classification else 'mse' + } + + return self.template_manager.render('pytorch/files/train.py.jinja2', context) + + def _generate_dataset_script(self, nodes: List[Dict[str, Any]]) -> str: + """Generate dataset script using template""" + input_shape = self._extract_input_shape(nodes) + + context = { + 'data_type': 'image', # Default to image + 'input_shape': input_shape, + 'input_channels': input_shape[1] if len(input_shape) > 1 else 3, + 'input_height': input_shape[2] if len(input_shape) > 2 else 224, + 'input_width': input_shape[3] if len(input_shape) > 3 else 224, + 'channel_type': 'RGB' if input_shape[1] == 3 else 'Grayscale' if input_shape[1] == 1 else f'{input_shape[1]}-channel' + } + + return self.template_manager.render('pytorch/files/dataset.py.jinja2', context) + + def _generate_config_file(self, nodes: List[Dict[str, Any]]) -> str: + """Generate config file using template""" + input_shape = self._extract_input_shape(nodes) + + # Count layers + layer_count = sum( + 1 for n in nodes + if get_node_type(n) not in ('input', 'output', 'dataloader') + ) + + # Determine complexity and hyperparameters + if layer_count > 20: + batch_size = 16 + learning_rate = 1e-4 + epochs = 100 + complexity = "Deep" + elif layer_count > 10: + batch_size = 32 + learning_rate = 1e-3 + epochs = 50 + complexity = "Medium" + else: + batch_size = 64 + learning_rate = 1e-3 + epochs = 30 + complexity = "Shallow" + + # Check for attention layers + has_attention = any(get_node_type(n) in ('self_attention', 'attention') for n in nodes) + if has_attention: + learning_rate = learning_rate * 0.1 + batch_size = max(8, batch_size // 2) + + context = { + 'batch_size': batch_size, + 'learning_rate': learning_rate, + 'num_epochs': epochs, + 'input_shape': list(input_shape), + 'complexity': complexity, + 'layer_count': layer_count, + 'has_attention': has_attention + } + + return self.template_manager.render('pytorch/files/config.py.jinja2', context) diff --git a/project/block_manager/services/codegen/tensorflow_group_generator.py b/project/block_manager/services/codegen/tensorflow_group_generator.py new file mode 100644 index 0000000..60dffb6 --- /dev/null +++ b/project/block_manager/services/codegen/tensorflow_group_generator.py @@ -0,0 +1,323 @@ +""" +TensorFlow Group Block Generator +Generates TensorFlow keras.Model code for group block definitions. +""" + +from typing import Dict, List, Any, Optional +from .group_block_generator import GroupBlockGenerator +from ..nodes.base import Framework, LayerCodeSpec +from ..nodes.registry import get_node_definition +from .base import get_node_type, get_node_config + + +class TensorFlowGroupBlockGenerator(GroupBlockGenerator): + """TensorFlow-specific group block code generation""" + + def __init__(self): + super().__init__(Framework.TENSORFLOW) + + def generate_group_block_spec( + self, + group_definition: Dict[str, Any], + node_id: str, + instance_config: Optional[Dict[str, Any]] = None + ) -> LayerCodeSpec: + """ + Generate LayerCodeSpec for TensorFlow group block instance. + + Args: + group_definition: The GroupBlockDefinition dict + node_id: The node ID of the group instance in main graph + instance_config: Optional per-instance config overrides + + Returns: + LayerCodeSpec for this group instance + """ + class_name = self._sanitize_class_name(group_definition['name']) + sanitized_id = node_id.replace('-', '_') + layer_var_name = f"{sanitized_id}_{class_name}" + + # Extract init params from internal structure + init_params = self._extract_init_params(group_definition, instance_config) + + # Build template context + internal_structure = group_definition.get('internal_structure', {}) + port_mappings = internal_structure.get('portMappings', []) + input_ports, output_ports = self._parse_port_mappings(port_mappings) + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var_name, + node_type='group', + node_id=node_id, + init_params=init_params, + config_params=instance_config or {}, + template_context={ + 'group_definition_id': group_definition['id'], + 'has_multi_output': len(output_ports) > 1, + 'num_outputs': len(output_ports), + 'num_inputs': len(input_ports) + } + ) + + def generate_group_class_code( + self, + group_definition: Dict[str, Any] + ) -> str: + """ + Generate TensorFlow group block class using template. + + Args: + group_definition: The GroupBlockDefinition dict + + Returns: + Rendered class code as string + """ + internal_structure = group_definition.get('internal_structure', {}) + internal_nodes = internal_structure.get('nodes', []) + internal_edges = internal_structure.get('edges', []) + port_mappings = internal_structure.get('portMappings', []) + + # Sort internal nodes + sorted_nodes = self._topologically_sort_internal_nodes( + internal_nodes, internal_edges + ) + + # Parse ports + input_ports, output_ports = self._parse_port_mappings(port_mappings) + + # Generate LayerCodeSpecs for internal nodes + internal_specs = self._generate_internal_node_specs(sorted_nodes) + + # Build edge map + edge_map = self._build_edge_map(internal_edges) + + # Generate forward pass lines + forward_lines = self._generate_call_method( + sorted_nodes, + internal_specs, + edge_map, + input_ports, + output_ports + ) + + # Build template context + context = self._build_tensorflow_template_context( + group_definition, + sorted_nodes, + internal_specs, + input_ports, + output_ports, + forward_lines + ) + + # Render template + template_path = 'tensorflow/layers/group_block.py.jinja2' + return self.template_manager.render(template_path, context) + + def _generate_internal_node_specs( + self, + internal_nodes: List[Dict[str, Any]] + ) -> List[LayerCodeSpec]: + """ + Generate code specs for each internal node. + + Args: + internal_nodes: Sorted list of internal nodes + + Returns: + List of LayerCodeSpec for processable internal nodes + """ + specs = [] + + for node in internal_nodes: + node_type = get_node_type(node) + + # Skip special nodes + if node_type in ('input', 'output', 'dataloader'): + continue + + node_id = node['id'] + config = get_node_config(node) + + # Check if this is a nested group + if node_type == 'group': + # For Phase 1-2, we can defer nested groups + # TODO: Implement nested group support + continue + + # Get node definition from registry + node_def = get_node_definition(node_type, Framework.TENSORFLOW) + if node_def: + spec = node_def.get_tensorflow_code_spec( + node_id=node_id, + config=config, + input_shape=None, # Shape inference deferred + output_shape=None + ) + specs.append(spec) + + return specs + + def _generate_call_method( + self, + sorted_nodes: List[Dict[str, Any]], + internal_specs: List[LayerCodeSpec], + edge_map: Dict[str, List[str]], + input_ports: List[Dict], + output_ports: List[Dict] + ) -> List[str]: + """ + Generate call method lines for internal graph. + + Args: + sorted_nodes: Topologically sorted internal nodes + internal_specs: LayerCodeSpecs for internal nodes + edge_map: Map of node_id -> [incoming_node_ids] + input_ports: Input port mappings + output_ports: Output port mappings + + Returns: + List of call method code lines + """ + call_lines = [] + var_map = {} + spec_map = {spec.node_id: spec for spec in internal_specs} + + # Map input ports to their external parameter names + input_node_map = { + port['internalNodeId']: port + for port in input_ports + } + + # Track which nodes are outputs (need to preserve their variables) + output_node_ids = {port['internalNodeId'] for port in output_ports} + + for node in sorted_nodes: + node_id = node['id'] + node_type = get_node_type(node) + + # Handle input nodes + if node_type == 'input': + # Map input node to its external parameter + if node_id in input_node_map: + # For TensorFlow, inputs come as list or single value + if len(input_ports) == 1: + var_map[node_id] = 'inputs' + else: + # Multi-input: unpack from inputs list + idx = list(input_node_map.keys()).index(node_id) + var_map[node_id] = f'inputs[{idx}]' + continue + + # Skip output and dataloader nodes + if node_type in ('output', 'dataloader'): + continue + + # Get the spec for this node + spec = spec_map.get(node_id) + if not spec: + continue + + # Determine input variable + incoming = edge_map.get(node_id, []) + if not incoming: + input_var = 'inputs' + elif len(incoming) == 1: + input_var = var_map.get(incoming[0], 'inputs') + else: + # Multiple inputs (add, concat nodes) + input_vars = [var_map.get(src, 'inputs') for src in incoming] + input_var = f"[{', '.join(input_vars)}]" + + # Generate output variable + output_var = f"x_{node_id.replace('-', '_')}" + + # Generate layer call with training parameter + if node_type in ('add', 'concat'): + # Merge nodes don't use training parameter + if node_type == 'concat': + axis = spec.template_context.get('axis', -1) + call_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var}, axis={axis})" + ) + else: + call_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var})" + ) + elif node_type in ('dropout', 'batchnorm'): + # Layers that use training parameter + call_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var}, training=training)" + ) + else: + # Standard layer call + call_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var})" + ) + + var_map[node_id] = output_var + + return call_lines + + def _build_tensorflow_template_context( + self, + group_definition: Dict[str, Any], + sorted_nodes: List[Dict[str, Any]], + internal_specs: List[LayerCodeSpec], + input_ports: List[Dict], + output_ports: List[Dict], + call_lines: List[str] + ) -> Dict[str, Any]: + """ + Build comprehensive template context for TensorFlow group block. + + Args: + group_definition: The group definition + sorted_nodes: Sorted internal nodes + internal_specs: LayerCodeSpecs for internal nodes + input_ports: Input port mappings + output_ports: Output port mappings + call_lines: Generated call method code lines + + Returns: + Complete template context dict + """ + base_context = self._build_template_context( + group_definition, sorted_nodes, input_ports, output_ports + ) + + # Build output variables + output_node_ids = [port['internalNodeId'] for port in output_ports] + output_vars = [] + for node_id in output_node_ids: + # Use the variable name from call method + var_name = f"x_{node_id.replace('-', '_')}" + output_vars.append(var_name) + + return { + **base_context, + 'internal_specs': internal_specs, + 'call_lines': call_lines, + 'output_vars': output_vars, + 'init_params': [] # No custom init params for Phase 1-2 + } + + def _extract_init_params( + self, + group_definition: Dict[str, Any], + instance_config: Optional[Dict[str, Any]] + ) -> Dict[str, Any]: + """ + Extract initialization parameters for group block instance. + + Args: + group_definition: The group definition + instance_config: Per-instance config overrides + + Returns: + Dict of init parameters (empty for Phase 1-2) + """ + # For Phase 1-2, group blocks don't have instance-level init params + # All config is baked into the class definition + return {} diff --git a/project/block_manager/services/codegen/tensorflow_orchestrator.py b/project/block_manager/services/codegen/tensorflow_orchestrator.py new file mode 100644 index 0000000..53f54c0 --- /dev/null +++ b/project/block_manager/services/codegen/tensorflow_orchestrator.py @@ -0,0 +1,719 @@ +""" +TensorFlow Code Generation Orchestrator +Coordinates the generation of complete TensorFlow/Keras project files +""" + +from typing import List, Dict, Any, Optional, Tuple, Set +from collections import defaultdict +import json + +from .base import topological_sort, get_input_variable, get_node_type, get_node_config +from ..nodes.registry import get_node_definition +from ..nodes.base import Framework, LayerCodeSpec +from ..nodes.templates.manager import TemplateManager + + +class UnsupportedNodeTypeError(Exception): + """Raised when a node type is not supported""" + pass + + +class TensorFlowCodeOrchestrator: + """ + Orchestrator for TensorFlow/Keras code generation. + Delegates code generation to individual node classes and assembles the final output. + """ + + def __init__(self): + self.template_manager = TemplateManager() + + def generate( + self, + nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]], + project_name: str = "GeneratedModel", + group_definitions: Optional[List[Dict[str, Any]]] = None + ) -> Tuple[Dict[str, str], List[Exception]]: + """ + Generate complete TensorFlow/Keras project files. + + Args: + nodes: List of node definitions from the frontend + edges: List of edge definitions + project_name: Name for the generated model class + group_definitions: Optional list of group block definitions + + Returns: + Tuple of (files dict, errors list) + files dict contains: {'model': str, 'train': str, 'dataset': str, 'config': str} + """ + errors = [] + + try: + # Initialize group block generator if needed + group_generator = None + if group_definitions: + from .tensorflow_group_generator import TensorFlowGroupBlockGenerator + group_generator = TensorFlowGroupBlockGenerator() + + # Sort nodes topologically + sorted_nodes = topological_sort(nodes, edges) + + # Build edge map for quick lookups + edge_map = self._build_edge_map(edges) + + # Generate code specifications for each node + code_specs, spec_errors = self._generate_code_specs( + sorted_nodes, edge_map, group_generator, group_definitions + ) + errors.extend(spec_errors) + + # Generate code specs for internal layers in group blocks + if group_definitions: + internal_specs, internal_errors = self._generate_internal_layer_specs( + group_definitions + ) + code_specs.extend(internal_specs) + errors.extend(internal_errors) + + # Render layer classes from templates (includes internal layers) + layer_classes = self._render_layer_classes(code_specs) + + # Generate group block class definitions + group_classes = "" + if group_generator and group_definitions: + group_classes = self._generate_group_block_classes( + group_definitions, group_generator + ) + + # Combine regular layers + group classes + all_classes = layer_classes + if group_classes: + all_classes += "\n\n" + group_classes + + # Generate model class definition + model_definition = self._generate_model_definition( + project_name, + code_specs, + sorted_nodes, + edges, + edge_map + ) + + # Generate test code + input_shape = self._extract_input_shape(nodes) + test_code = self._generate_test_code(project_name, input_shape) + + # Render complete model file + model_code = self._render_model_file( + project_name, + all_classes, + model_definition, + test_code + ) + + # Generate training script + train_code = self._generate_training_script(project_name, nodes) + + # Generate dataset script + dataset_code = self._generate_dataset_script(nodes) + + # Generate config file + config_code = self._generate_config_file(nodes) + + return { + 'model': model_code, + 'train': train_code, + 'dataset': dataset_code, + 'config': config_code + }, errors + + except Exception as e: + errors.append(e) + return {}, errors + + def _build_edge_map(self, edges: List[Dict[str, Any]]) -> Dict[str, List[str]]: + """Build a map of node_id -> list of incoming node_ids""" + edge_map = defaultdict(list) + for edge in edges: + target = edge.get('target') + source = edge.get('source') + if target and source: + edge_map[target].append(source) + return dict(edge_map) + + def _build_outgoing_edge_map(self, edges: List[Dict[str, Any]]) -> Dict[str, List[str]]: + """ + Build a map of outgoing edges from each node. + + Returns: + Dict mapping source_node_id -> [target_node_ids] + """ + outgoing_map = defaultdict(list) + for edge in edges: + source = edge.get('source') + target = edge.get('target') + if source and target: + outgoing_map[source].append(target) + return dict(outgoing_map) + + def _generate_code_specs( + self, + sorted_nodes: List[Dict[str, Any]], + edge_map: Dict[str, List[str]], + group_generator=None, + group_definitions: Optional[List[Dict[str, Any]]] = None + ) -> Tuple[List[LayerCodeSpec], List[Exception]]: + """Generate code specifications for all nodes including group blocks""" + code_specs = [] + errors = [] + + # Build group definition lookup + group_def_map = {} + if group_definitions: + group_def_map = {gd['id']: gd for gd in group_definitions} + + processable_nodes = [ + n for n in sorted_nodes + if get_node_type(n) not in ('input', 'dataloader', 'output') + ] + + for node in processable_nodes: + try: + node_id = node['id'] + node_type = get_node_type(node) + config = get_node_config(node) + + # Handle group blocks + if node_type == 'group': + if not group_generator: + raise UnsupportedNodeTypeError( + f"Group node {node_id} found but no group_definitions provided" + ) + + group_def_id = node.get('data', {}).get('groupDefinitionId') + group_def = group_def_map.get(group_def_id) + + if not group_def: + raise ValueError( + f"Group definition {group_def_id} not found for node {node_id}" + ) + + # Generate spec for this group instance + code_spec = group_generator.generate_group_block_spec( + group_definition=group_def, + node_id=node_id, + instance_config=config + ) + code_specs.append(code_spec) + + else: + # Regular node + node_def = get_node_definition(node_type, Framework.TENSORFLOW) + + if not node_def: + raise UnsupportedNodeTypeError( + f"Node type '{node_type}' (id: {node_id}) is not supported for TensorFlow" + ) + + code_spec = node_def.get_tensorflow_code_spec( + node_id=node_id, + config=config, + input_shape=None, + output_shape=None + ) + + code_specs.append(code_spec) + + except Exception as e: + errors.append(e) + + return code_specs, errors + + def _generate_internal_layer_specs( + self, + group_definitions: List[Dict[str, Any]] + ) -> Tuple[List[LayerCodeSpec], List[Exception]]: + """ + Generate LayerCodeSpecs for all unique internal layers used in group blocks. + This ensures internal layer classes are defined before group blocks use them. + + Args: + group_definitions: List of group block definitions + + Returns: + Tuple of (list of internal layer specs, list of errors) + """ + internal_specs = [] + errors = [] + seen_node_types = set() + + for group_def in group_definitions: + internal_structure = group_def.get('internal_structure', {}) + internal_nodes = internal_structure.get('nodes', []) + + for node in internal_nodes: + node_type = get_node_type(node) + + # Skip special nodes + if node_type in ('input', 'output', 'dataloader', 'group'): + continue + + # Only generate each node type once + if node_type in seen_node_types: + continue + + seen_node_types.add(node_type) + + try: + node_id = node['id'] + config = get_node_config(node) + + # Get node definition from registry + node_def = get_node_definition(node_type, Framework.TENSORFLOW) + + if not node_def: + raise UnsupportedNodeTypeError( + f"Internal node type '{node_type}' not supported in group block" + ) + + # Generate code specification + code_spec = node_def.get_tensorflow_code_spec( + node_id=node_id, + config=config, + input_shape=None, + output_shape=None + ) + + internal_specs.append(code_spec) + + except Exception as e: + errors.append(e) + + return internal_specs, errors + + def _render_layer_classes(self, code_specs: List[LayerCodeSpec]) -> str: + """Render all unique layer class definitions""" + unique_classes = {} + + for spec in code_specs: + if spec.node_type not in unique_classes: + try: + template_path = spec.get_template_path(Framework.TENSORFLOW) + # Merge class_name, init_params, and template_context for rendering + context = { + 'class_name': spec.class_name, + **spec.init_params, + **spec.template_context + } + rendered = self.template_manager.render(template_path, context) + unique_classes[spec.node_type] = rendered + except Exception: + pass + + return '\n\n'.join(unique_classes.values()) + + def _generate_group_block_classes( + self, + group_definitions: List[Dict[str, Any]], + group_generator + ) -> str: + """ + Generate class definitions for all group blocks. + + Args: + group_definitions: List of group block definitions + group_generator: TensorFlowGroupBlockGenerator instance + + Returns: + String containing all group block class definitions + """ + # Detect dependency order for nested groups + ordered_definitions = self._order_group_definitions(group_definitions) + + class_codes = [] + for group_def in ordered_definitions: + try: + class_code = group_generator.generate_group_class_code(group_def) + class_codes.append(class_code) + except Exception as e: + # Log error but continue with other groups + print(f"Error generating group {group_def.get('name')}: {e}") + + return '\n\n'.join(class_codes) + + def _order_group_definitions( + self, + group_definitions: List[Dict[str, Any]] + ) -> List[Dict[str, Any]]: + """ + Order group definitions so nested groups are defined before their parents. + Uses topological sort based on group dependencies. + + Args: + group_definitions: List of group block definitions + + Returns: + Topologically sorted list of group definitions + """ + # Build dependency graph + graph = {gd['id']: [] for gd in group_definitions} + + for group_def in group_definitions: + internal_nodes = group_def.get('internal_structure', {}).get('nodes', []) + for node in internal_nodes: + if node.get('data', {}).get('blockType') == 'group': + nested_group_id = node.get('data', {}).get('groupDefinitionId') + if nested_group_id in graph: + # This group depends on the nested group + graph[group_def['id']].append(nested_group_id) + + # Topological sort (simple implementation) + visited = set() + result = [] + + def visit(gd_id): + if gd_id in visited: + return + visited.add(gd_id) + for dep in graph.get(gd_id, []): + visit(dep) + result.append(gd_id) + + for gd in group_definitions: + visit(gd['id']) + + # Map back to group definitions + gd_map = {gd['id']: gd for gd in group_definitions} + return [gd_map[gd_id] for gd_id in result if gd_id in gd_map] + + def _generate_model_definition( + self, + project_name: str, + code_specs: List[LayerCodeSpec], + sorted_nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]], + edge_map: Dict[str, List[str]] + ) -> str: + """Generate the main model class definition""" + # Generate layer initializations + layer_inits = [] + for spec in code_specs: + params_str = ', '.join( + f"{k}={repr(v)}" for k, v in spec.init_params.items() + ) + layer_inits.append( + f"self.{spec.layer_variable_name} = {spec.class_name}({params_str})" + ) + + # Generate forward pass logic + forward_lines, _ = self._generate_forward_pass( + sorted_nodes, + edges, + edge_map, + code_specs + ) + + model_class = f'''class {project_name}(keras.Model): + """ + TensorFlow/Keras Model for {project_name} + + This model is auto-generated from the VisionForge architecture. + """ + + def __init__(self): + super({project_name}, self).__init__() + #========================== + #Layer Initializations: + #========================== +{chr(10).join(" " + line for line in layer_inits)} + + def call(self, inputs, training=None): + """ + Forward pass through the model. + + Args: + inputs: Input tensor (NHWC format) + training: Whether in training mode + + Returns: + Output tensor after processing through the model + """ + x = inputs +{chr(10).join(" " + line for line in forward_lines)} + return x +''' + return model_class + + def _needs_named_variable( + self, + node_id: str, + node_type: str, + outgoing_edge_map: Dict[str, List[str]], + num_outputs: int + ) -> bool: + """ + Determine if a node's output requires a named variable. + + A named variable is needed when: + 1. Node has multiple outgoing edges (output used multiple times - skip connections) + 2. Node has multiple outputs (e.g., group blocks) + 3. Node is a merge operation (add, concat) - for readability + + Args: + node_id: Node identifier + node_type: Type of node (e.g., 'conv2d', 'add') + outgoing_edge_map: Map of node_id -> list of target node IDs + num_outputs: Number of outputs for this node + + Returns: + True if a named variable should be created + """ + # Multi-output nodes always need named variables + if num_outputs > 1: + return True + + # Nodes with multiple outgoing edges need named variables (skip connections) + outgoing_edges = outgoing_edge_map.get(node_id, []) + if len(outgoing_edges) > 1: + return True + + # Merge operations benefit from named variables for readability + if node_type in ('add', 'concat'): + return True + + return False + + def _generate_forward_pass( + self, + sorted_nodes: List[Dict[str, Any]], + edges: List[Dict[str, Any]], + edge_map: Dict[str, List[str]], + code_specs: List[LayerCodeSpec] + ) -> Tuple[List[str], Set[str]]: + """ + Generate forward pass logic with optimized variable usage. + + Only creates named variables when needed: + - Skip connections (nodes with multiple outgoing edges) + - Multi-output nodes (group blocks) + - Merge operations (add, concat) + + Otherwise reuses 'x' variable for memory efficiency. + + Returns: + Tuple of (forward pass lines, set of skip connection var names) + """ + forward_lines = [] + var_map = {} + skip_connections = set() + spec_map = {spec.node_id: spec for spec in code_specs} + + # Build outgoing edge map to detect skip connections + outgoing_edge_map = self._build_outgoing_edge_map(edges) + + processable_nodes = [ + n for n in sorted_nodes + if get_node_type(n) not in ('output',) + ] + + for node in processable_nodes: + node_id = node['id'] + node_type = get_node_type(node) + + if node_type in ('input', 'dataloader'): + var_map[node_id] = 'x' + continue + + incoming = edge_map.get(node_id, []) + input_var = get_input_variable(incoming, var_map) + + spec = spec_map.get(node_id) + if not spec: + continue + + # Handle group blocks with multiple outputs + if node_type == 'group' and spec.template_context.get('has_multi_output'): + num_outputs = spec.template_context.get('num_outputs', 1) + output_vars = [ + f"{node_id.replace('-', '_')}_out{i}" + for i in range(num_outputs) + ] + forward_lines.append( + f"{', '.join(output_vars)} = self.{spec.layer_variable_name}({input_var}, training=training)" + ) + # Map first output as primary variable for this node + var_map[node_id] = output_vars[0] + skip_connections.add(output_vars[0]) + + else: + # Determine if this node needs a named variable + num_outputs = 1 + needs_named_var = self._needs_named_variable( + node_id, + node_type, + outgoing_edge_map, + num_outputs + ) + + if needs_named_var: + # Create named variable for skip connections, merge ops, etc. + output_var = f"x_{node_id.replace('-', '_')}" + forward_lines.append( + f"{output_var} = self.{spec.layer_variable_name}({input_var}, training=training)" + ) + var_map[node_id] = output_var + + # Track if this is a skip connection source + if len(outgoing_edge_map.get(node_id, [])) > 1: + skip_connections.add(output_var) + + else: + # Reuse 'x' variable for linear chains (memory efficient) + forward_lines.append( + f"x = self.{spec.layer_variable_name}({input_var}, training=training)" + ) + var_map[node_id] = 'x' + + # Ensure final output is assigned to 'x' + if processable_nodes: + last_node_id = processable_nodes[-1]['id'] + last_var = var_map.get(last_node_id, 'x') + if last_var != 'x': + forward_lines.append(f"x = {last_var}") + + return forward_lines, skip_connections + + def _extract_input_shape(self, nodes: List[Dict[str, Any]]) -> Tuple[int, ...]: + """Extract input shape from input node""" + input_node = next((n for n in nodes if get_node_type(n) == 'input'), None) + + if input_node: + config = get_node_config(input_node) + shape_str = config.get('shape', '[1, 224, 224, 3]') + try: + shape = json.loads(shape_str) if isinstance(shape_str, str) else shape_str + if isinstance(shape, list): + return tuple(shape) + except (ValueError, TypeError): + pass + + return (1, 224, 224, 3) # NHWC format default + + def _generate_test_code(self, project_name: str, input_shape: Tuple[int, ...]) -> str: + """Generate test code for model validation""" + return f'''if __name__ == "__main__": + # Test the model with random input + model = {project_name}() + test_input = tf.random.normal({input_shape}) + print(f"Input shape: {{test_input.shape}}") + output = model(test_input, training=False) + print(f"Output shape: {{output.shape}}") + print(f"Model has {{model.count_params():,}} parameters") +''' + + def _render_model_file( + self, + project_name: str, + layer_classes: str, + model_definition: str, + test_code: str + ) -> str: + """Render the complete model.py file""" + return f'''""" +Generated TensorFlow/Keras Model +Architecture: {project_name} +Generated by VisionForge + +This file contains the model architecture with separate layer classes. +Each layer is implemented as a reusable class for clarity and maintainability. +""" + +import tensorflow as tf +from tensorflow import keras +from tensorflow.keras import layers +from typing import List, Tuple, Optional + + +#========================== +#Layer Definitions: +#========================== +{layer_classes} + +{model_definition} + +{test_code} +''' + + def _generate_training_script(self, project_name: str, nodes: List[Dict[str, Any]]) -> str: + """Generate training script using template""" + has_softmax = any(get_node_type(n) == 'softmax' for n in nodes) + is_classification = has_softmax + + context = { + 'project_name': project_name, + 'model_class_name': project_name, + 'task_type': 'classification' if is_classification else 'regression', + 'is_classification': is_classification, + 'loss_function': 'keras.losses.SparseCategoricalCrossentropy()' if is_classification else 'keras.losses.MeanSquaredError()', + 'metric_name': 'accuracy' if is_classification else 'mse' + } + + return self.template_manager.render('tensorflow/files/train.py.jinja2', context) + + def _generate_dataset_script(self, nodes: List[Dict[str, Any]]) -> str: + """Generate dataset script using template""" + input_shape = self._extract_input_shape(nodes) + + context = { + 'data_type': 'image', + 'input_shape': input_shape, + 'input_height': input_shape[1] if len(input_shape) > 1 else 224, + 'input_width': input_shape[2] if len(input_shape) > 2 else 224, + 'input_channels': input_shape[3] if len(input_shape) > 3 else 3, + 'channel_type': 'RGB' if input_shape[3] == 3 else 'Grayscale' if input_shape[3] == 1 else f'{input_shape[3]}-channel' + } + + return self.template_manager.render('tensorflow/files/dataset.py.jinja2', context) + + def _generate_config_file(self, nodes: List[Dict[str, Any]]) -> str: + """Generate config file using template""" + input_shape = self._extract_input_shape(nodes) + + layer_count = sum( + 1 for n in nodes + if get_node_type(n) not in ('input', 'output', 'dataloader') + ) + + if layer_count > 20: + batch_size = 16 + learning_rate = 1e-4 + epochs = 100 + complexity = "Deep" + elif layer_count > 10: + batch_size = 32 + learning_rate = 1e-3 + epochs = 50 + complexity = "Medium" + else: + batch_size = 64 + learning_rate = 1e-3 + epochs = 30 + complexity = "Shallow" + + has_attention = any(get_node_type(n) in ('self_attention', 'attention') for n in nodes) + if has_attention: + learning_rate = learning_rate * 0.1 + batch_size = max(8, batch_size // 2) + + context = { + 'batch_size': batch_size, + 'learning_rate': learning_rate, + 'num_epochs': epochs, + 'input_shape': list(input_shape), + 'complexity': complexity, + 'layer_count': layer_count, + 'has_attention': has_attention + } + + return self.template_manager.render('tensorflow/files/config.py.jinja2', context) diff --git a/project/block_manager/services/enhanced_pytorch_codegen.py b/project/block_manager/services/enhanced_pytorch_codegen.py index 6ea9ad3..6b7a8e9 100644 --- a/project/block_manager/services/enhanced_pytorch_codegen.py +++ b/project/block_manager/services/enhanced_pytorch_codegen.py @@ -2,6 +2,9 @@ from collections import deque import json +# New template-based code generation +from .codegen.pytorch_orchestrator import PyTorchCodeOrchestrator + # ============================================ # Custom Exception Classes # ============================================ @@ -1553,7 +1556,16 @@ def generate_pytorch_code( code files (with keys: 'model', 'train', 'dataset', 'config') and a list of shape errors encountered during generation. """ - # Generate layer class definitions + # Delegate to new template-based orchestrator + orchestrator = PyTorchCodeOrchestrator() + return orchestrator.generate(nodes, edges, project_name, group_definitions) + + # ==================== LEGACY CODE (PRESERVED FOR REFERENCE) ==================== + # The code below is no longer executed but preserved for reference. + # All code generation now happens through the PyTorchCodeOrchestrator. + # ================================================================================ + + # OLD: Generate layer class definitions layer_classes = ClassDefinitionGenerator.create_required_node_classes(nodes, group_definitions) # Generate layer initializations diff --git a/project/block_manager/services/nodes/base.py b/project/block_manager/services/nodes/base.py index e31f947..a6a9947 100644 --- a/project/block_manager/services/nodes/base.py +++ b/project/block_manager/services/nodes/base.py @@ -6,6 +6,7 @@ from abc import ABC, abstractmethod from typing import Dict, List, Optional, Any, Tuple from enum import Enum +from dataclasses import dataclass, field class Framework(str, Enum): @@ -14,6 +15,27 @@ class Framework(str, Enum): TENSORFLOW = "tensorflow" +@dataclass +class LayerCodeSpec: + """ + Structured specification for generating layer code. + Contains all information needed to render a layer's code template. + """ + class_name: str + layer_variable_name: str + node_type: str + node_id: str + init_params: Dict[str, Any] = field(default_factory=dict) + config_params: Dict[str, Any] = field(default_factory=dict) + input_shape_info: Dict[str, Any] = field(default_factory=dict) + output_shape_info: Dict[str, Any] = field(default_factory=dict) + template_context: Dict[str, Any] = field(default_factory=dict) + + def get_template_path(self, framework: Framework) -> str: + """Get the template path for this layer type""" + return f"{framework.value}/layers/{self.node_type}.py.jinja2" + + class TensorShape: """Represents tensor shape with dimensions and description""" @@ -278,6 +300,52 @@ def to_dict(self) -> Dict[str, Any]: "configSchema": [field.to_dict() for field in self.config_schema] } + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """ + Generate specification for PyTorch code generation. + + Args: + node_id: Unique identifier for this node instance + config: Configuration parameters for this node + input_shape: Input tensor shape + output_shape: Output tensor shape + + Returns: + LayerCodeSpec with all information needed to render PyTorch code + """ + raise NotImplementedError( + f"{self.__class__.__name__} must implement get_pytorch_code_spec()" + ) + + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """ + Generate specification for TensorFlow code generation. + + Args: + node_id: Unique identifier for this node instance + config: Configuration parameters for this node + input_shape: Input tensor shape (NHWC format) + output_shape: Output tensor shape (NHWC format) + + Returns: + LayerCodeSpec with all information needed to render TensorFlow code + """ + raise NotImplementedError( + f"{self.__class__.__name__} must implement get_tensorflow_code_spec()" + ) + class SourceNodeDefinition(NodeDefinition): """Base class for input/source nodes that don't receive connections""" diff --git a/project/block_manager/services/nodes/pytorch/add.py b/project/block_manager/services/nodes/pytorch/add.py index 356870b..cb55c43 100644 --- a/project/block_manager/services/nodes/pytorch/add.py +++ b/project/block_manager/services/nodes/pytorch/add.py @@ -1,7 +1,7 @@ """PyTorch Add Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class AddNode(NodeDefinition): @@ -52,3 +52,27 @@ def validate_incoming_connection( def allows_multiple_inputs(self) -> bool: """Add nodes accept multiple input connections""" return True + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Add layer""" + sanitized_id = node_id.replace('-', '_') + class_name = 'AddBlock' + layer_var = f'{sanitized_id}_AddBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='add', + node_id=node_id, + init_params={}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={} + ) diff --git a/project/block_manager/services/nodes/pytorch/attention.py b/project/block_manager/services/nodes/pytorch/attention.py new file mode 100644 index 0000000..181efa5 --- /dev/null +++ b/project/block_manager/services/nodes/pytorch/attention.py @@ -0,0 +1,118 @@ +"""PyTorch Multi-Head Attention Node Definition""" + +from typing import Dict, List, Optional, Any +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec + + +class AttentionNode(NodeDefinition): + """Multi-Head Attention layer""" + + @property + def metadata(self) -> NodeMetadata: + return NodeMetadata( + type="attention", + label="Multi-Head Attention", + category="advanced", + color="var(--color-accent)", + icon="Zap", + description="Multi-head self-attention mechanism", + framework=Framework.PYTORCH + ) + + @property + def config_schema(self) -> List[ConfigField]: + return [ + ConfigField( + name="embed_dim", + label="Embedding Dimension", + type="number", + default=512, + min=1, + description="Total dimension of the model" + ), + ConfigField( + name="num_heads", + label="Number of Heads", + type="number", + default=8, + min=1, + description="Number of attention heads" + ), + ConfigField( + name="dropout", + label="Dropout", + type="number", + default=0.0, + min=0.0, + max=1.0, + description="Dropout probability" + ), + ConfigField( + name="bias", + label="Use Bias", + type="boolean", + default=True, + description="Whether to use bias in linear projections" + ) + ] + + def compute_output_shape( + self, + input_shape: Optional[TensorShape], + config: Dict[str, Any] + ) -> Optional[TensorShape]: + # Multi-head attention preserves shape + if input_shape: + return TensorShape( + dims=input_shape.dims, + description="Attention output" + ) + return None + + def validate_incoming_connection( + self, + source_node_type: str, + source_output_shape: Optional[TensorShape], + target_config: Dict[str, Any] + ) -> Optional[str]: + # Attention accepts any input shape (typically 3D: [batch, seq_len, embed_dim]) + return None + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Multi-Head Attention layer""" + embed_dim = config.get('embed_dim', 512) + num_heads = config.get('num_heads', 8) + dropout = config.get('dropout', 0.0) + bias = config.get('bias', True) + + sanitized_id = node_id.replace('-', '_') + class_name = 'MultiHeadAttentionLayer' + layer_var = f'{sanitized_id}_MultiHeadAttentionLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='attention', + node_id=node_id, + init_params={ + 'embed_dim': embed_dim, + 'num_heads': num_heads, + 'dropout': dropout, + 'bias': bias + }, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={ + 'embed_dim': embed_dim, + 'num_heads': num_heads, + 'dropout': dropout, + 'bias': bias + } + ) diff --git a/project/block_manager/services/nodes/pytorch/batchnorm2d.py b/project/block_manager/services/nodes/pytorch/batchnorm2d.py index 786dc9c..0ab1f9f 100644 --- a/project/block_manager/services/nodes/pytorch/batchnorm2d.py +++ b/project/block_manager/services/nodes/pytorch/batchnorm2d.py @@ -1,7 +1,7 @@ """PyTorch BatchNorm2D Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class BatchNorm2DNode(NodeDefinition): @@ -95,3 +95,42 @@ def validate_incoming_connection( 4, "[batch, channels, height, width]" ) + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for BatchNorm2D layer""" + num_features = config.get('num_features') or (input_shape.dims[1] if input_shape and len(input_shape.dims) >= 2 else 64) + eps = config.get('eps', 1e-5) + momentum = config.get('momentum', 0.1) + affine = config.get('affine', True) + + sanitized_id = node_id.replace('-', '_') + class_name = 'BatchNormBlock' + layer_var = f'{sanitized_id}_BatchNormBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='batchnorm', + node_id=node_id, + init_params={ + 'num_features': num_features, + 'eps': eps, + 'momentum': momentum, + 'affine': affine + }, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={ + 'num_features': num_features, + 'eps': eps, + 'momentum': momentum, + 'affine': affine + } + ) diff --git a/project/block_manager/services/nodes/pytorch/concat.py b/project/block_manager/services/nodes/pytorch/concat.py index 04f4fa7..cd24c4e 100644 --- a/project/block_manager/services/nodes/pytorch/concat.py +++ b/project/block_manager/services/nodes/pytorch/concat.py @@ -1,7 +1,7 @@ """PyTorch Concat Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class ConcatNode(NodeDefinition): @@ -63,3 +63,29 @@ def validate_incoming_connection( def allows_multiple_inputs(self) -> bool: """Concat nodes accept multiple input connections""" return True + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Concat layer""" + dim = config.get('dim', 1) + + sanitized_id = node_id.replace('-', '_') + class_name = 'ConcatBlock' + layer_var = f'{sanitized_id}_ConcatBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='concat', + node_id=node_id, + init_params={}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={'dim': dim} + ) diff --git a/project/block_manager/services/nodes/pytorch/conv2d.py b/project/block_manager/services/nodes/pytorch/conv2d.py index 806be44..db5c51b 100644 --- a/project/block_manager/services/nodes/pytorch/conv2d.py +++ b/project/block_manager/services/nodes/pytorch/conv2d.py @@ -1,7 +1,7 @@ """PyTorch Conv2D Layer Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class Conv2DNode(NodeDefinition): @@ -110,3 +110,58 @@ def validate_incoming_connection( 4, "[batch, channels, height, width]" ) + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Conv2D layer""" + # Extract parameters from config + out_channels = config.get('out_channels', 64) + kernel_size = config.get('kernel_size', 3) + stride = config.get('stride', 1) + padding = config.get('padding', 0) + dilation = config.get('dilation', 1) + + # Extract input channels from shape + in_channels = input_shape.dims[1] if input_shape and len(input_shape.dims) >= 2 else 3 + + # Generate unique class and variable names + sanitized_id = node_id.replace('-', '_') + class_name = 'Conv2DBlock' + layer_var = f'{sanitized_id}_Conv2DBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='conv2d', + node_id=node_id, + init_params={ + 'in_channels': in_channels, + 'out_channels': out_channels, + 'kernel_size': kernel_size, + 'stride': stride, + 'padding': padding, + 'dilation': dilation + }, + config_params=config, + input_shape_info={ + 'in_channels': in_channels, + 'dims': input_shape.dims if input_shape else [] + }, + output_shape_info={ + 'out_channels': out_channels, + 'dims': output_shape.dims if output_shape else [] + }, + template_context={ + 'in_channels': in_channels, + 'out_channels': out_channels, + 'kernel_size': kernel_size, + 'stride': stride, + 'padding': padding, + 'dilation': dilation + } + ) diff --git a/project/block_manager/services/nodes/pytorch/dropout.py b/project/block_manager/services/nodes/pytorch/dropout.py index b6b547c..df6de6f 100644 --- a/project/block_manager/services/nodes/pytorch/dropout.py +++ b/project/block_manager/services/nodes/pytorch/dropout.py @@ -1,7 +1,7 @@ """PyTorch Dropout Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class DropoutNode(NodeDefinition): @@ -61,3 +61,29 @@ def validate_incoming_connection( ) -> Optional[str]: # Dropout accepts any input shape return None + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Dropout layer""" + p = config.get('p', 0.5) + + sanitized_id = node_id.replace('-', '_') + class_name = 'DropoutLayer' + layer_var = f'{sanitized_id}_DropoutLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='dropout', + node_id=node_id, + init_params={'p': p}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={'p': p} + ) diff --git a/project/block_manager/services/nodes/pytorch/flatten.py b/project/block_manager/services/nodes/pytorch/flatten.py index 90808c6..9f5a41b 100644 --- a/project/block_manager/services/nodes/pytorch/flatten.py +++ b/project/block_manager/services/nodes/pytorch/flatten.py @@ -1,7 +1,7 @@ """PyTorch Flatten Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class FlattenNode(NodeDefinition): @@ -77,5 +77,39 @@ def validate_incoming_connection( # Validate that input has at least 2 dimensions if source_output_shape and len(source_output_shape.dims) < 2: return "Flatten requires input with at least 2 dimensions" - + return None + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Flatten layer""" + start_dim = config.get('start_dim', 1) + + out_features = output_shape.dims[1] if output_shape and len(output_shape.dims) >= 2 else None + + sanitized_id = node_id.replace('-', '_') + class_name = 'FlattenLayer' + layer_var = f'{sanitized_id}_FlattenLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='flatten', + node_id=node_id, + init_params={'start_dim': start_dim}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={ + 'dims': output_shape.dims if output_shape else [], + 'out_features': out_features + }, + template_context={ + 'start_dim': start_dim, + 'out_features': out_features + } + ) diff --git a/project/block_manager/services/nodes/pytorch/linear.py b/project/block_manager/services/nodes/pytorch/linear.py index d8b167e..8a701b2 100644 --- a/project/block_manager/services/nodes/pytorch/linear.py +++ b/project/block_manager/services/nodes/pytorch/linear.py @@ -1,7 +1,7 @@ """PyTorch Linear Layer Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class LinearNode(NodeDefinition): @@ -75,3 +75,46 @@ def validate_incoming_connection( 2, "[batch, features]" ) + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Linear layer""" + out_features = config.get('out_features', 128) + bias = config.get('bias', True) + + in_features = input_shape.dims[1] if input_shape and len(input_shape.dims) >= 2 else 512 + + sanitized_id = node_id.replace('-', '_') + class_name = 'LinearLayer' + layer_var = f'{sanitized_id}_LinearLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='linear', + node_id=node_id, + init_params={ + 'in_features': in_features, + 'out_features': out_features, + 'bias': bias + }, + config_params=config, + input_shape_info={ + 'in_features': in_features, + 'dims': input_shape.dims if input_shape else [] + }, + output_shape_info={ + 'out_features': out_features, + 'dims': output_shape.dims if output_shape else [] + }, + template_context={ + 'in_features': in_features, + 'out_features': out_features, + 'bias': bias + } + ) diff --git a/project/block_manager/services/nodes/pytorch/maxpool2d.py b/project/block_manager/services/nodes/pytorch/maxpool2d.py index 2d77b9a..5f456ed 100644 --- a/project/block_manager/services/nodes/pytorch/maxpool2d.py +++ b/project/block_manager/services/nodes/pytorch/maxpool2d.py @@ -1,7 +1,7 @@ """PyTorch MaxPool2D Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class MaxPool2DNode(NodeDefinition): @@ -99,3 +99,39 @@ def validate_incoming_connection( 4, "[batch, channels, height, width]" ) + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for MaxPool2D layer""" + kernel_size = config.get('kernel_size', 2) + stride = config.get('stride', 2) + padding = config.get('padding', 0) + + sanitized_id = node_id.replace('-', '_') + class_name = 'MaxPoolBlock' + layer_var = f'{sanitized_id}_MaxPoolBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='maxpool', + node_id=node_id, + init_params={ + 'kernel_size': kernel_size, + 'stride': stride, + 'padding': padding + }, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={ + 'kernel_size': kernel_size, + 'stride': stride, + 'padding': padding + } + ) diff --git a/project/block_manager/services/nodes/pytorch/relu.py b/project/block_manager/services/nodes/pytorch/relu.py new file mode 100644 index 0000000..c9e5564 --- /dev/null +++ b/project/block_manager/services/nodes/pytorch/relu.py @@ -0,0 +1,80 @@ +"""PyTorch ReLU Activation Node Definition""" + +from typing import Dict, List, Optional, Any +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec + + +class ReLUNode(NodeDefinition): + """ReLU activation function layer""" + + @property + def metadata(self) -> NodeMetadata: + return NodeMetadata( + type="relu", + label="ReLU", + category="basic", + color="var(--color-primary)", + icon="Zap", + description="ReLU activation function", + framework=Framework.PYTORCH + ) + + @property + def config_schema(self) -> List[ConfigField]: + return [ + ConfigField( + name="inplace", + label="In-place", + type="boolean", + default=False, + description="Perform operation in-place" + ) + ] + + def compute_output_shape( + self, + input_shape: Optional[TensorShape], + config: Dict[str, Any] + ) -> Optional[TensorShape]: + # ReLU preserves shape + if input_shape: + return TensorShape( + dims=input_shape.dims, + description="ReLU activated" + ) + return None + + def validate_incoming_connection( + self, + source_node_type: str, + source_output_shape: Optional[TensorShape], + target_config: Dict[str, Any] + ) -> Optional[str]: + # ReLU accepts any input shape + return None + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for ReLU layer""" + inplace = config.get('inplace', False) + + sanitized_id = node_id.replace('-', '_') + class_name = 'ReLULayer' + layer_var = f'{sanitized_id}_ReLULayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='relu', + node_id=node_id, + init_params={'inplace': inplace}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={'inplace': inplace} + ) diff --git a/project/block_manager/services/nodes/pytorch/softmax.py b/project/block_manager/services/nodes/pytorch/softmax.py new file mode 100644 index 0000000..c2b6eca --- /dev/null +++ b/project/block_manager/services/nodes/pytorch/softmax.py @@ -0,0 +1,80 @@ +"""PyTorch Softmax Activation Node Definition""" + +from typing import Dict, List, Optional, Any +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec + + +class SoftmaxNode(NodeDefinition): + """Softmax activation function layer""" + + @property + def metadata(self) -> NodeMetadata: + return NodeMetadata( + type="softmax", + label="Softmax", + category="basic", + color="var(--color-primary)", + icon="Activity", + description="Softmax activation function", + framework=Framework.PYTORCH + ) + + @property + def config_schema(self) -> List[ConfigField]: + return [ + ConfigField( + name="dim", + label="Dimension", + type="number", + default=1, + description="Dimension along which softmax will be computed" + ) + ] + + def compute_output_shape( + self, + input_shape: Optional[TensorShape], + config: Dict[str, Any] + ) -> Optional[TensorShape]: + # Softmax preserves shape + if input_shape: + return TensorShape( + dims=input_shape.dims, + description="Softmax probabilities" + ) + return None + + def validate_incoming_connection( + self, + source_node_type: str, + source_output_shape: Optional[TensorShape], + target_config: Dict[str, Any] + ) -> Optional[str]: + # Softmax accepts any input shape + return None + + def get_pytorch_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate PyTorch code specification for Softmax layer""" + dim = config.get('dim', 1) + + sanitized_id = node_id.replace('-', '_') + class_name = 'SoftmaxLayer' + layer_var = f'{sanitized_id}_SoftmaxLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='softmax', + node_id=node_id, + init_params={'dim': dim}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={'dim': dim} + ) diff --git a/project/block_manager/services/nodes/templates/__init__.py b/project/block_manager/services/nodes/templates/__init__.py new file mode 100644 index 0000000..95c1af2 --- /dev/null +++ b/project/block_manager/services/nodes/templates/__init__.py @@ -0,0 +1,5 @@ +"""Template management for code generation""" + +from .manager import TemplateManager + +__all__ = ['TemplateManager'] diff --git a/project/block_manager/services/nodes/templates/manager.py b/project/block_manager/services/nodes/templates/manager.py new file mode 100644 index 0000000..74514ca --- /dev/null +++ b/project/block_manager/services/nodes/templates/manager.py @@ -0,0 +1,82 @@ +""" +Template Manager for Code Generation +Handles loading and caching of Jinja2 templates +""" + +from jinja2 import Environment, PackageLoader, Template +from typing import Dict, Optional + + +class TemplateManager: + """ + Manages Jinja2 template loading and caching for code generation. + Singleton pattern ensures templates are loaded once and reused. + """ + + _instance: Optional['TemplateManager'] = None + _env: Optional[Environment] = None + _template_cache: Dict[str, Template] = {} + + def __new__(cls): + if cls._instance is None: + cls._instance = super().__new__(cls) + return cls._instance + + @classmethod + def get_environment(cls) -> Environment: + """Get or create the Jinja2 environment""" + if cls._env is None: + cls._env = Environment( + loader=PackageLoader( + 'block_manager.services.nodes', + 'templates' + ), + autoescape=False, # Generating Python code, not HTML + trim_blocks=True, + lstrip_blocks=True, + keep_trailing_newline=True + ) + + # Add custom filters if needed + cls._env.filters['repr'] = repr + + return cls._env + + @classmethod + def get_template(cls, template_path: str) -> Template: + """ + Get a template by path, with caching. + + Args: + template_path: Path to template relative to templates directory + e.g., "pytorch/layers/conv2d.py.jinja2" + + Returns: + Compiled Jinja2 template + """ + if template_path not in cls._template_cache: + env = cls.get_environment() + cls._template_cache[template_path] = env.get_template(template_path) + + return cls._template_cache[template_path] + + @classmethod + def render(cls, template_path: str, context: Dict) -> str: + """ + Render a template with the given context. + + Args: + template_path: Path to template + context: Template variables + + Returns: + Rendered template as string + """ + template = cls.get_template(template_path) + return template.render(**context) + + @classmethod + def clear_cache(cls): + """Clear the template cache (useful for testing)""" + cls._template_cache.clear() + cls._env = None diff --git a/project/block_manager/services/nodes/templates/pytorch/files/config.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/files/config.py.jinja2 new file mode 100644 index 0000000..91ab274 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/files/config.py.jinja2 @@ -0,0 +1,29 @@ +""" +Configuration File +Generated by VisionForge +Architecture Complexity: {{ complexity }} ({{ layer_count }} layers) +""" + +# Training Configuration +BATCH_SIZE = {{ batch_size }} # Adjusted for {{ complexity.lower() }} network +LEARNING_RATE = {{ learning_rate }} # {% if has_attention %}Reduced for attention layers{% else %}Standard for architecture{% endif %} +NUM_EPOCHS = {{ num_epochs }} +WEIGHT_DECAY = 1e-4 + +# Model Configuration (NCHW format: batch, channels, height, width) +INPUT_SHAPE = {{ input_shape }} + +# Data Configuration +DATA_DIR = './data' +NUM_WORKERS = 0 # Set to 0 for debugging, increase for faster data loading + +# Device Configuration +DEVICE = 'cuda' # Change to 'cpu' if no GPU available + +# Logging Configuration +LOG_INTERVAL = 10 # Print every N batches +SAVE_INTERVAL = 5 # Save checkpoint every N epochs + +# Architecture Info +NUM_LAYERS = {{ layer_count }} +HAS_ATTENTION = {{ has_attention }} diff --git a/project/block_manager/services/nodes/templates/pytorch/files/dataset.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/files/dataset.py.jinja2 new file mode 100644 index 0000000..2a9ddc6 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/files/dataset.py.jinja2 @@ -0,0 +1,92 @@ +""" +Custom Dataset Class +Generated by VisionForge +Data Type: {{ data_type|capitalize }}{% if input_channels %} ({{ input_channels }} channels, {{ input_height }}x{{ input_width }}){% endif %} +""" + +import torch +from torch.utils.data import Dataset +import numpy as np +from pathlib import Path +from typing import Tuple, Optional +{% if data_type == 'image' %} +from PIL import Image +import torchvision.transforms as transforms +{% endif %} + + +class CustomDataset(Dataset): + """ + Custom dataset for loading and preprocessing {{ data_type }} data. +{% if data_type == 'image' and input_channels %} + + Expected input: Images with shape [{{ input_channels }}, {{ input_height }}, {{ input_width }}] + Channels: {{ input_channels }} ({{ channel_type }}) +{% endif %} + """ + + def __init__(self, data_dir: str = './data', train: bool = True): + """ + Initialize the dataset. + + Args: + data_dir: Directory containing the data + train: Whether this is training data + """ + self.data_dir = Path(data_dir) + self.train = train + +{% if data_type == 'image' %} + # Define transforms + self.transform = transforms.Compose([ + transforms.Resize(({{ input_height }}, {{ input_width }})), +{% if input_channels == 1 %} + transforms.Grayscale(), +{% endif %} + transforms.ToTensor(), + transforms.Normalize( + mean=[0.5] * {{ input_channels }}, + std=[0.5] * {{ input_channels }} + ) + ]) +{% endif %} + + # TODO: Load your data here + # Example: self.images = list((self.data_dir / ('train' if train else 'val')).glob('*.jpg')) + # For now, we'll use dummy data + self.data = [] + self.labels = [] + + def __len__(self) -> int: + """Return the total number of samples""" + return len(self.data) if self.data else 100 # Dummy size + + def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Get a single sample. + + Args: + idx: Index of the sample + + Returns: + Tuple of (data tensor, label tensor) + """ + # TODO: Replace with actual data loading +{% if data_type == 'image' %} + # Example: + # image = Image.open(self.images[idx]) + # image = self.transform(image) + # label = self.labels[idx] + + # Dummy data for now + image = torch.randn({{ input_channels }}, {{ input_height }}, {{ input_width }}) + label = torch.randint(0, 10, (1,)).item() + + return image, label +{% else %} + # Dummy data for now + data = torch.randn(*{{ input_shape[1:] }}) + label = torch.randint(0, 10, (1,)).item() + + return data, label +{% endif %} diff --git a/project/block_manager/services/nodes/templates/pytorch/files/model.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/files/model.py.jinja2 new file mode 100644 index 0000000..ac250a7 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/files/model.py.jinja2 @@ -0,0 +1,58 @@ +""" +Generated PyTorch Model +Architecture: {{ project_name }} +Generated by VisionForge + +This file contains the model architecture with separate layer classes. +Each layer is implemented as a reusable class for clarity and maintainability. +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +from typing import List, Tuple, Optional + + +#========================== +#Layer Definitions: +#========================== +{% for layer_class in layer_classes %} +{{ layer_class }} + +{% endfor %} + +class {{ model_class_name }}(nn.Module): + """ + PyTorch Model for {{ project_name }} + + This model is auto-generated from the VisionForge architecture. + """ + + def __init__(self): + super({{ model_class_name }}, self).__init__() + #========================== + #Layer Initializations: + #========================== +{% for init_line in layer_initializations %} + {{ init_line }} +{% endfor %} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the model. + + Args: + x: Input tensor + + Returns: + Output tensor after processing through the model + """ +{% for forward_line in forward_pass_lines %} + {{ forward_line }} +{% endfor %} + return x + + +{% if test_code %} +{{ test_code }} +{% endif %} diff --git a/project/block_manager/services/nodes/templates/pytorch/files/train.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/files/train.py.jinja2 new file mode 100644 index 0000000..3f07dee --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/files/train.py.jinja2 @@ -0,0 +1,193 @@ +""" +Training Script for {{ project_name }} +Generated by VisionForge +Architecture Type: {{ task_type|capitalize }} +""" + +import torch +import torch.nn as nn +import torch.optim as optim +from torch.utils.data import DataLoader +from pathlib import Path +from typing import Tuple, Dict +import time + +from model import {{ model_class_name }} +from dataset import CustomDataset +from config import * + + +def train_epoch( + model: nn.Module, + dataloader: DataLoader, + criterion: nn.Module, + optimizer: optim.Optimizer, + device: torch.device +) -> Tuple[float, float]: + """ + Train for one epoch. + + Args: + model: The model to train + dataloader: Training data loader + criterion: Loss function + optimizer: Optimizer + device: Device to train on + + Returns: + Tuple of (average loss, metric value) + """ + model.train() + total_loss = 0.0 + correct = 0 + total = 0 + + for batch_idx, (data, target) in enumerate(dataloader): + data, target = data.to(device), target.to(device) + + optimizer.zero_grad() + output = model(data) + loss = criterion(output, target) + + loss.backward() + optimizer.step() + + total_loss += loss.item() + + # Calculate metric +{% if is_classification %} + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() +{% else %} + # For regression tasks, metric tracking can be added here +{% endif %} + total += target.size(0) + + avg_loss = total_loss / len(dataloader) +{% if is_classification %} + metric = 100. * correct / total +{% else %} + metric = avg_loss +{% endif %} + + return avg_loss, metric + + +def validate_epoch( + model: nn.Module, + dataloader: DataLoader, + criterion: nn.Module, + device: torch.device +) -> Tuple[float, float]: + """ + Validate for one epoch. + + Args: + model: The model to validate + dataloader: Validation data loader + criterion: Loss function + device: Device to validate on + + Returns: + Tuple of (average loss, metric value) + """ + model.eval() + total_loss = 0.0 + correct = 0 + total = 0 + + with torch.no_grad(): + for data, target in dataloader: + data, target = data.to(device), target.to(device) + output = model(data) + loss = criterion(output, target) + + total_loss += loss.item() + + # Calculate metric +{% if is_classification %} + pred = output.argmax(dim=1, keepdim=True) + correct += pred.eq(target.view_as(pred)).sum().item() +{% else %} + # Metric calculation for regression +{% endif %} + total += target.size(0) + + avg_loss = total_loss / len(dataloader) +{% if is_classification %} + metric = 100. * correct / total +{% else %} + metric = avg_loss +{% endif %} + + return avg_loss, metric + + +def main(): + """Main training function""" + # Set device + device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + print(f"Using device: {device}") + + # Create model + model = {{ model_class_name }}().to(device) + print(f"Model created with {sum(p.numel() for p in model.parameters()):,} parameters") + + # Create datasets + train_dataset = CustomDataset(train=True) + val_dataset = CustomDataset(train=False) + + train_loader = DataLoader( + train_dataset, + batch_size=BATCH_SIZE, + shuffle=True, + num_workers=0 + ) + + val_loader = DataLoader( + val_dataset, + batch_size=BATCH_SIZE, + shuffle=False, + num_workers=0 + ) + + # Setup training + criterion = {{ loss_function }} + optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) + scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5) + + # Training loop + best_val_loss = float('inf') + + for epoch in range(NUM_EPOCHS): + start_time = time.time() + + # Train + train_loss, train_metric = train_epoch(model, train_loader, criterion, optimizer, device) + + # Validate + val_loss, val_metric = validate_epoch(model, val_loader, criterion, device) + + # Update learning rate + scheduler.step(val_loss) + + epoch_time = time.time() - start_time + + # Print progress + print(f"Epoch {epoch+1}/{NUM_EPOCHS} | " + f"Time: {epoch_time:.2f}s | " + f"Train Loss: {train_loss:.4f} | Train {{ metric_name.upper() }}: {train_metric:.2f} | " + f"Val Loss: {val_loss:.4f} | Val {{ metric_name.upper() }}: {val_metric:.2f}") + + # Save best model + if val_loss < best_val_loss: + best_val_loss = val_loss + torch.save(model.state_dict(), 'best_model.pth') + print(f" → New best model saved (val_loss: {val_loss:.4f})") + + print("\nTraining completed!") + print(f"Best validation loss: {best_val_loss:.4f}") + + +if __name__ == "__main__": + main() diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/_base_layer.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/_base_layer.py.jinja2 new file mode 100644 index 0000000..e965837 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/_base_layer.py.jinja2 @@ -0,0 +1,39 @@ +class {{ class_name }}(nn.Module): + """ + {{ docstring }} + {% if parameters %} + + Parameters: + {% for param_name, param_desc in parameters.items() %} + - {{ param_desc }} + {% endfor %} + {% endif %} + {% if shape_info %} + + Shape: + - Input: {{ shape_info.input }} + - Output: {{ shape_info.output }} + {% endif %} + """ + + def __init__(self{% if init_signature %}, {{ init_signature }}{% endif %}): + """Initialize the {{ layer_type }} layer.""" + super({{ class_name }}, self).__init__() + {% block init_body %} + # Layer initialization + pass + {% endblock %} + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the {{ layer_type }} layer. + + Args: + x: Input tensor{% if input_shape_desc %} of shape {{ input_shape_desc }}{% endif %} + + Returns: + Output tensor{% if output_shape_desc %} of shape {{ output_shape_desc }}{% endif %} + """ + {% block forward_body %} + return x + {% endblock %} diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/add.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/add.py.jinja2 new file mode 100644 index 0000000..caa2f38 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/add.py.jinja2 @@ -0,0 +1,25 @@ +class {{ class_name }}(nn.Module): + """ + Addition Block + + Element-wise addition block for tensors of same shape. + + Parameters: + - tensor_list: list of input tensors of same shape + """ + + def __init__(self): + """Initialize the addition layer (no learnable parameters).""" + super({{ class_name }}, self).__init__() + + def forward(self, tensor_list: List[torch.Tensor]) -> torch.Tensor: + """ + Forward pass through the addition layer. + + Args: + tensor_list: List of input tensors + + Returns: + Element-wise sum of all input tensors + """ + return torch.stack(tensor_list).sum(dim=0) diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/attention.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/attention.py.jinja2 new file mode 100644 index 0000000..694f84c --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/attention.py.jinja2 @@ -0,0 +1,40 @@ +class {{ class_name }}(nn.Module): + """ + Multi-Head Self-Attention Layer + + Applies multi-head self-attention mechanism to the input. + Allows the model to jointly attend to information from different representation subspaces. + + Parameters: + - Embedding dimension: {{ embed_dim }} + - Number of heads: {{ num_heads }} + - Dropout: {{ dropout }} + + Shape: + - Input: [batch_size, seq_len, {{ embed_dim }}] + - Output: [batch_size, seq_len, {{ embed_dim }}] + """ + + def __init__(self, embed_dim: int = {{ embed_dim }}, num_heads: int = {{ num_heads }}, dropout: float = {{ dropout }}, bias: bool = True): + """Initialize the multi-head attention layer.""" + super({{ class_name }}, self).__init__() + self.attention = nn.MultiheadAttention( + embed_dim=embed_dim, + num_heads=num_heads, + dropout=dropout, + batch_first=True, + bias=bias + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the attention layer. + + Args: + x: Input tensor of shape [batch, seq_len, {{ embed_dim }}] + + Returns: + Output tensor after applying multi-head attention + """ + x, _ = self.attention(x, x, x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/batchnorm.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/batchnorm.py.jinja2 new file mode 100644 index 0000000..9b651c7 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/batchnorm.py.jinja2 @@ -0,0 +1,40 @@ +class {{ class_name }}(nn.Module): + """ + Batch Normalization Layer + + Normalizes the input over a mini-batch for each feature channel. + Helps stabilize and accelerate training. + + Parameters: + - Number of features: {{ num_features }} + - Epsilon: {{ eps }} + - Momentum: {{ momentum }} + - Learnable parameters: {{ affine }} + + Shape: + - Input: [batch_size, {{ num_features }}, H, W] + - Output: [batch_size, {{ num_features }}, H, W] + """ + + def __init__(self, num_features: int, eps: float = {{ eps }}, momentum: float = {{ momentum }}, affine: bool = {{ affine }}): + """Initialize the batch normalization layer.""" + super({{ class_name }}, self).__init__() + self.bn = nn.BatchNorm2d( + num_features, + eps=eps, + momentum=momentum, + affine=affine + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the batch normalization layer. + + Args: + x: Input tensor of shape [batch, {{ num_features }}, H, W] + + Returns: + Normalized output tensor of same shape + """ + x = self.bn(x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/concat.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/concat.py.jinja2 new file mode 100644 index 0000000..e948624 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/concat.py.jinja2 @@ -0,0 +1,31 @@ +class {{ class_name }}(nn.Module): + """ + Concatenation Layer + + Concatenates multiple tensors along a specified dimension. + Commonly used to merge feature maps from different paths in the network. + + Parameters: + - Concatenation dimension: {{ dim }} + + Shape: + - Input: List of tensors with compatible shapes + - Output: Concatenated tensor along dimension {{ dim }} + """ + + def __init__(self): + """Initialize the concatenation layer (no learnable parameters).""" + super({{ class_name }}, self).__init__() + + def forward(self, tensor_list: List[torch.Tensor], concat_dim: int = {{ dim }}): + """ + Forward pass through the concatenation layer. + + Args: + tensor_list: List of input tensors + concat_dim: Dimension along which to concatenate + + Returns: + Concatenated tensor along dimension concat_dim + """ + return torch.cat(tensor_list, dim=concat_dim) diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/conv2d.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/conv2d.py.jinja2 new file mode 100644 index 0000000..1ede786 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/conv2d.py.jinja2 @@ -0,0 +1,43 @@ +class {{ class_name }}(nn.Module): + """ + 2D Convolutional Layer + + Applies a 2D convolution over an input signal composed of several input channels. + + Parameters: + - Input channels: {{ in_channels }} + - Output channels: {{ out_channels }} + - Kernel size: {{ kernel_size }}x{{ kernel_size }} + - Stride: {{ stride }} + - Padding: {{ padding }} + - Dilation: {{ dilation }} + + Shape: + - Input: [batch_size, {{ in_channels }}, H, W] + - Output: [batch_size, {{ out_channels }}, H_out, W_out] + """ + + def __init__(self, in_channels: int, out_channels: int, kernel_size={{ kernel_size }}, stride={{ stride }}, padding={{ padding }}, dilation={{ dilation }}): + """Initialize the convolutional layer.""" + super({{ class_name }}, self).__init__() + self.conv = nn.Conv2d( + in_channels, + out_channels, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the convolutional layer. + + Args: + x: Input tensor of shape [batch, {{ in_channels }}, H, W] + + Returns: + Output tensor of shape [batch, {{ out_channels }}, H_out, W_out] + """ + x = self.conv(x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/dropout.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/dropout.py.jinja2 new file mode 100644 index 0000000..7d2ce18 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/dropout.py.jinja2 @@ -0,0 +1,32 @@ +class {{ class_name }}(nn.Module): + """ + Dropout Regularization Layer + + Randomly zeroes some elements of the input tensor with probability p during training. + Helps prevent overfitting. + + Parameters: + - Dropout probability: {{ p }} + + Shape: + - Input: [batch_size, *] (any shape) + - Output: [batch_size, *] (same shape as input) + """ + + def __init__(self, p: float = {{ p }}, inplace: bool = False): + """Initialize the dropout layer.""" + super({{ class_name }}, self).__init__() + self.dropout = nn.Dropout(p=p, inplace=inplace) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the dropout layer. + + Args: + x: Input tensor + + Returns: + Output tensor with dropout applied during training + """ + x = self.dropout(x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/flatten.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/flatten.py.jinja2 new file mode 100644 index 0000000..d23041f --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/flatten.py.jinja2 @@ -0,0 +1,29 @@ +class {{ class_name }}(nn.Module): + """ + Flatten Layer + + Flattens a contiguous range of dimensions into a tensor. + Commonly used to transition from convolutional layers to fully connected layers. + + Shape: + - Input: [batch_size, C, H, W] + - Output: [batch_size, C*H*W]{% if out_features %} = [batch_size, {{ out_features }}]{% endif %} + """ + + def __init__(self, start_dim: int = 1): + """Initialize the flatten layer.""" + super({{ class_name }}, self).__init__() + self.flatten = nn.Flatten(start_dim=start_dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the flatten layer. + + Args: + x: Input tensor of shape [batch, C, H, W] + + Returns: + Output tensor of shape [batch, C*H*W] + """ + x = self.flatten(x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/group_block.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/group_block.py.jinja2 new file mode 100644 index 0000000..18561a0 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/group_block.py.jinja2 @@ -0,0 +1,65 @@ +class {{ class_name }}(nn.Module): + """ + {{ group_name }}{% if description %}: {{ description }}{% endif %} + + Custom group block with internal layers. +{% if has_multi_input %} + + Inputs: +{% for port in input_ports %} + - {{ port.externalPortLabel }}: {{ port.semantic }} (port: {{ port.externalPortId }}) +{% endfor %} +{% endif %} +{% if has_multi_output %} + + Outputs: +{% for port in output_ports %} + - {{ port.externalPortLabel }}: {{ port.semantic }} (port: {{ port.externalPortId }}) +{% endfor %} +{% endif %} + """ + + def __init__(self{% for param in init_params %}, {{ param.name }}={{ param.default }}{% endfor %}): + """Initialize {{ group_name }} with internal layers.""" + super({{ class_name }}, self).__init__() + + # Internal layers +{% for spec in internal_specs %} + self.{{ spec.layer_variable_name }} = {{ spec.class_name }}( +{% for key, value in spec.init_params.items() %} + {{ key }}={{ value }}{{ "," if not loop.last else "" }} +{% endfor %} + ) +{% endfor %} + + def forward(self, {{ input_params }}) -> {% if has_multi_output %}tuple{% else %}torch.Tensor{% endif %}: + """ + Forward pass through {{ group_name }}. +{% if has_multi_input %} + + Args: +{% for port in input_ports %} + {{ port.externalPortLabel.lower().replace(' ', '_') }}: Input tensor for {{ port.semantic }} +{% endfor %} +{% else %} + + Args: + x: Input tensor +{% endif %} + + Returns: +{% if has_multi_output %} + Tuple of output tensors: ({{ output_vars | join(', ') }}) +{% else %} + Output tensor +{% endif %} + """ +{% for line in forward_lines %} + {{ line }} +{% endfor %} + +{% if has_multi_output %} + return ({{ output_vars | join(', ') }}) +{% else %} + return {{ output_vars[0] if output_vars else 'x' }} +{% endif %} diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/linear.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/linear.py.jinja2 new file mode 100644 index 0000000..f88b1af --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/linear.py.jinja2 @@ -0,0 +1,33 @@ +class {{ class_name }}(nn.Module): + """ + Fully Connected Linear Layer + + Applies a linear transformation to the incoming data: y = xA^T + b + + Parameters: + - Input features: {{ in_features }} + - Output features: {{ out_features }} + - Bias: {{ bias }} + + Shape: + - Input: [batch_size, {{ in_features }}] + - Output: [batch_size, {{ out_features }}] + """ + + def __init__(self, in_features: int, out_features: int, bias: bool = {{ bias }}): + """Initialize the linear layer.""" + super({{ class_name }}, self).__init__() + self.linear = nn.Linear(in_features, out_features, bias=bias) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the linear layer. + + Args: + x: Input tensor of shape [batch, {{ in_features }}] + + Returns: + Output tensor of shape [batch, {{ out_features }}] + """ + x = self.linear(x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/maxpool.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/maxpool.py.jinja2 new file mode 100644 index 0000000..97709a4 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/maxpool.py.jinja2 @@ -0,0 +1,38 @@ +class {{ class_name }}(nn.Module): + """ + 2D Max Pooling Layer + + Applies a 2D max pooling over an input signal. + Reduces spatial dimensions while preserving channel count. + + Parameters: + - Kernel size: {{ kernel_size }}x{{ kernel_size }} + - Stride: {{ stride }} + - Padding: {{ padding }} + + Shape: + - Input: [batch_size, C, H, W] + - Output: [batch_size, C, H/stride, W/stride] + """ + + def __init__(self, kernel_size: int, stride: int, padding: int): + """Initialize the max pooling layer.""" + super({{ class_name }}, self).__init__() + self.pool = nn.MaxPool2d( + kernel_size=kernel_size, + stride=stride, + padding=padding + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the pooling layer. + + Args: + x: Input tensor of shape [batch, C, H, W] + + Returns: + Output tensor with reduced spatial dimensions + """ + x = self.pool(x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/relu.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/relu.py.jinja2 new file mode 100644 index 0000000..f8aadda --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/relu.py.jinja2 @@ -0,0 +1,29 @@ +class {{ class_name }}(nn.Module): + """ + ReLU Activation Layer + + Applies the rectified linear unit function element-wise: ReLU(x) = max(0, x) + Introduces non-linearity to the model. + + Shape: + - Input: [batch_size, *] (any shape) + - Output: [batch_size, *] (same shape as input) + """ + + def __init__(self, inplace: bool = False): + """Initialize the ReLU activation.""" + super({{ class_name }}, self).__init__() + self.relu = nn.ReLU(inplace=inplace) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the activation. + + Args: + x: Input tensor + + Returns: + Output tensor with ReLU applied element-wise + """ + x = self.relu(x) + return x diff --git a/project/block_manager/services/nodes/templates/pytorch/layers/softmax.py.jinja2 b/project/block_manager/services/nodes/templates/pytorch/layers/softmax.py.jinja2 new file mode 100644 index 0000000..db1a690 --- /dev/null +++ b/project/block_manager/services/nodes/templates/pytorch/layers/softmax.py.jinja2 @@ -0,0 +1,32 @@ +class {{ class_name }}(nn.Module): + """ + Softmax Activation Layer + + Applies the softmax function to normalize outputs into a probability distribution. + Commonly used in the final layer for classification tasks. + + Parameters: + - Dimension: {{ dim }} + + Shape: + - Input: [batch_size, num_classes] + - Output: [batch_size, num_classes] (sums to 1.0 along dimension {{ dim }}) + """ + + def __init__(self, dim: int = {{ dim }}): + """Initialize the softmax layer.""" + super({{ class_name }}, self).__init__() + self.softmax = nn.Softmax(dim=dim) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass through the softmax layer. + + Args: + x: Input tensor + + Returns: + Probability distribution over dimension {{ dim }} + """ + x = self.softmax(x) + return x diff --git a/project/block_manager/services/nodes/templates/tensorflow/files/config.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/files/config.py.jinja2 new file mode 100644 index 0000000..14a9573 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/files/config.py.jinja2 @@ -0,0 +1,21 @@ +""" +Configuration File +Generated by VisionForge +Architecture Complexity: {{ complexity }} ({{ layer_count }} layers) +Framework: TensorFlow/Keras +""" + +# Training Configuration +BATCH_SIZE = {{ batch_size }} # Adjusted for {{ complexity.lower() }} network +LEARNING_RATE = {{ learning_rate }} # {% if has_attention %}Reduced for attention layers{% else %}Standard for architecture{% endif %} +NUM_EPOCHS = {{ num_epochs }} + +# Model Configuration (NHWC format: batch, height, width, channels) +INPUT_SHAPE = {{ input_shape }} + +# Data Configuration +DATA_DIR = './data' + +# Architecture Info +NUM_LAYERS = {{ layer_count }} +HAS_ATTENTION = {{ has_attention }} diff --git a/project/block_manager/services/nodes/templates/tensorflow/files/dataset.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/files/dataset.py.jinja2 new file mode 100644 index 0000000..87d0528 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/files/dataset.py.jinja2 @@ -0,0 +1,81 @@ +""" +Dataset Creation for TensorFlow +Generated by VisionForge +Data Type: {{ data_type|capitalize }}{% if input_channels %} ({{ input_channels }} channels, {{ input_height }}x{{ input_width }}){% endif %} +""" + +import tensorflow as tf +import numpy as np +from pathlib import Path +from typing import Tuple + + +def create_dataset(train: bool = True, batch_size: int = 32): + """ + Create a TensorFlow dataset for loading and preprocessing {{ data_type }} data. + + Args: + train: Whether this is training data + batch_size: Batch size for the dataset + + Returns: + tf.data.Dataset instance +{% if data_type == 'image' and input_channels %} + + Expected input: Images with shape [{{ input_height }}, {{ input_width }}, {{ input_channels }}] (NHWC format) + Channels: {{ input_channels }} ({{ channel_type }}) +{% endif %} + """ + data_dir = Path('./data') + + # TODO: Load your actual data here + # Example for image data: + # train_dir = data_dir / ('train' if train else 'val') + # dataset = tf.keras.preprocessing.image_dataset_from_directory( + # train_dir, + # image_size=({{ input_height }}, {{ input_width }}), + # batch_size=batch_size, + # label_mode='int' + # ) + + # For now, create dummy data + num_samples = 1000 if train else 200 + + def data_generator(): + for _ in range(num_samples): +{% if data_type == 'image' %} + # Generate dummy image data (NHWC format) + image = tf.random.normal([{{ input_height }}, {{ input_width }}, {{ input_channels }}]) + label = tf.random.uniform([], minval=0, maxval=10, dtype=tf.int32) + yield image, label +{% else %} + # Generate dummy data + data = tf.random.normal({{ input_shape[1:] }}) + label = tf.random.uniform([], minval=0, maxval=10, dtype=tf.int32) + yield data, label +{% endif %} + + # Create dataset from generator + dataset = tf.data.Dataset.from_generator( + data_generator, +{% if data_type == 'image' %} + output_signature=( + tf.TensorSpec(shape=[{{ input_height }}, {{ input_width }}, {{ input_channels }}], dtype=tf.float32), + tf.TensorSpec(shape=[], dtype=tf.int32) + ) +{% else %} + output_signature=( + tf.TensorSpec(shape={{ input_shape[1:] }}, dtype=tf.float32), + tf.TensorSpec(shape=[], dtype=tf.int32) + ) +{% endif %} + ) + + # Shuffle and batch + if train: + dataset = dataset.shuffle(buffer_size=1000) + + dataset = dataset.batch(batch_size) + dataset = dataset.prefetch(tf.data.AUTOTUNE) + + return dataset diff --git a/project/block_manager/services/nodes/templates/tensorflow/files/model.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/files/model.py.jinja2 new file mode 100644 index 0000000..5ff0adf --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/files/model.py.jinja2 @@ -0,0 +1,59 @@ +""" +Generated TensorFlow/Keras Model +Architecture: {{ project_name }} +Generated by VisionForge + +This file contains the model architecture with separate layer classes. +Each layer is implemented as a reusable class for clarity and maintainability. +""" + +import tensorflow as tf +from tensorflow import keras +from tensorflow.keras import layers +from typing import List, Tuple, Optional + + +#========================== +#Layer Definitions: +#========================== +{% for layer_class in layer_classes %} +{{ layer_class }} + +{% endfor %} + +class {{ model_class_name }}(keras.Model): + """ + TensorFlow/Keras Model for {{ project_name }} + + This model is auto-generated from the VisionForge architecture. + """ + + def __init__(self): + super({{ model_class_name }}, self).__init__() + #========================== + #Layer Initializations: + #========================== +{% for init_line in layer_initializations %} + {{ init_line }} +{% endfor %} + + def call(self, inputs, training=None): + """ + Forward pass through the model. + + Args: + inputs: Input tensor (NHWC format) + training: Whether in training mode + + Returns: + Output tensor after processing through the model + """ +{% for forward_line in forward_pass_lines %} + {{ forward_line }} +{% endfor %} + return x + + +{% if test_code %} +{{ test_code }} +{% endif %} diff --git a/project/block_manager/services/nodes/templates/tensorflow/files/train.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/files/train.py.jinja2 new file mode 100644 index 0000000..b242120 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/files/train.py.jinja2 @@ -0,0 +1,124 @@ +""" +Training Script for {{ project_name }} +Generated by VisionForge +Architecture Type: {{ task_type|capitalize }} +Framework: TensorFlow/Keras +""" + +import tensorflow as tf +from tensorflow import keras +import numpy as np +from pathlib import Path +from typing import Tuple, Dict +import time + +from model import {{ model_class_name }} +from dataset import create_dataset +from config import * + + +def train_step(model: keras.Model, x_batch, y_batch, loss_fn, optimizer): + """ + Perform one training step. + + Args: + model: The model to train + x_batch: Input batch + y_batch: Target batch + loss_fn: Loss function + optimizer: Optimizer + + Returns: + Loss value for this batch + """ + with tf.GradientTape() as tape: + predictions = model(x_batch, training=True) + loss = loss_fn(y_batch, predictions) + + gradients = tape.gradient(loss, model.trainable_variables) + optimizer.apply_gradients(zip(gradients, model.trainable_variables)) + + return loss + + +def main(): + """Main training function""" + print(f"TensorFlow version: {tf.__version__}") + print(f"GPU Available: {tf.config.list_physical_devices('GPU')}") + + # Create model + model = {{ model_class_name }}() + print(f"Model created") + + # Create datasets + train_dataset = create_dataset(train=True, batch_size=BATCH_SIZE) + val_dataset = create_dataset(train=False, batch_size=BATCH_SIZE) + + # Setup training + loss_fn = {{ loss_function }} + optimizer = keras.optimizers.Adam(learning_rate=LEARNING_RATE) + +{% if is_classification %} + train_accuracy = keras.metrics.SparseCategoricalAccuracy(name='train_accuracy') + val_accuracy = keras.metrics.SparseCategoricalAccuracy(name='val_accuracy') +{% endif %} + + # Training loop + best_val_loss = float('inf') + + for epoch in range(NUM_EPOCHS): + start_time = time.time() + + # Training + train_loss_avg = keras.metrics.Mean() +{% if is_classification %} + train_accuracy.reset_states() +{% endif %} + + for x_batch, y_batch in train_dataset: + loss = train_step(model, x_batch, y_batch, loss_fn, optimizer) + train_loss_avg.update_state(loss) +{% if is_classification %} + train_accuracy.update_state(y_batch, model(x_batch, training=True)) +{% endif %} + + # Validation + val_loss_avg = keras.metrics.Mean() +{% if is_classification %} + val_accuracy.reset_states() +{% endif %} + + for x_batch, y_batch in val_dataset: + predictions = model(x_batch, training=False) + loss = loss_fn(y_batch, predictions) + val_loss_avg.update_state(loss) +{% if is_classification %} + val_accuracy.update_state(y_batch, predictions) +{% endif %} + + epoch_time = time.time() - start_time + + # Print progress + print(f"Epoch {epoch+1}/{NUM_EPOCHS} | " + f"Time: {epoch_time:.2f}s | " + f"Train Loss: {train_loss_avg.result():.4f} | " +{% if is_classification %} + f"Train Acc: {train_accuracy.result()*100:.2f}% | " + f"Val Loss: {val_loss_avg.result():.4f} | " + f"Val Acc: {val_accuracy.result()*100:.2f}%") +{% else %} + f"Val Loss: {val_loss_avg.result():.4f}") +{% endif %} + + # Save best model + if val_loss_avg.result() < best_val_loss: + best_val_loss = val_loss_avg.result() + model.save_weights('best_model.h5') + print(f" → New best model saved (val_loss: {best_val_loss:.4f})") + + print("\nTraining completed!") + print(f"Best validation loss: {best_val_loss:.4f}") + + +if __name__ == "__main__": + main() diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/_base_layer.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/_base_layer.py.jinja2 new file mode 100644 index 0000000..9a948c7 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/_base_layer.py.jinja2 @@ -0,0 +1,40 @@ +class {{ class_name }}(layers.Layer): + """ + {{ docstring }} + {% if parameters %} + + Parameters: + {% for param_name, param_desc in parameters.items() %} + - {{ param_desc }} + {% endfor %} + {% endif %} + {% if shape_info %} + + Shape: + - Input: {{ shape_info.input }} (NHWC format) + - Output: {{ shape_info.output }} (NHWC format) + {% endif %} + """ + + def __init__(self): + """Initialize the {{ layer_type }} layer.""" + super({{ class_name }}, self).__init__() + {% block init_body %} + # Layer initialization + pass + {% endblock %} + + def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: + """ + Forward pass through the {{ layer_type }} layer. + + Args: + inputs: Input tensor{% if input_shape_desc %} of shape {{ input_shape_desc }}{% endif %} + training: Whether in training mode + + Returns: + Output tensor{% if output_shape_desc %} of shape {{ output_shape_desc }}{% endif %} + """ + {% block forward_body %} + return inputs + {% endblock %} diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/add.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/add.py.jinja2 new file mode 100644 index 0000000..6015f92 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/add.py.jinja2 @@ -0,0 +1,31 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + Addition Layer + + Element-wise addition of multiple tensors. + + Parameters: + - Input tensors: List of tensors with same shape + + Shape: + - Input: List of tensors [batch_size, *] + - Output: [batch_size, *] (same shape as inputs) + """ + + def __init__(self): + """Initialize the addition layer.""" + super({{ class_name }}, self).__init__() + self.add = tf.keras.layers.Add() + + def call(self, inputs, training=None): + """ + Forward pass through the addition layer. + + Args: + inputs: List of input tensors + training: Whether in training mode + + Returns: + Element-wise sum of all input tensors + """ + return self.add(inputs) diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/batchnorm.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/batchnorm.py.jinja2 new file mode 100644 index 0000000..4e947f9 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/batchnorm.py.jinja2 @@ -0,0 +1,36 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + Batch Normalization Layer + + Normalizes the input over a mini-batch. + Helps stabilize and accelerate training. + + Parameters: + - Epsilon: {{ epsilon }} + - Momentum: {{ momentum }} + + Shape: + - Input: [batch_size, H, W, C] (NHWC format) + - Output: [batch_size, H, W, C] + """ + + def __init__(self, epsilon={{ epsilon }}, momentum={{ momentum }}): + """Initialize the batch normalization layer.""" + super({{ class_name }}, self).__init__() + self.bn = tf.keras.layers.BatchNormalization( + epsilon=epsilon, + momentum=momentum + ) + + def call(self, inputs, training=None): + """ + Forward pass through the batch normalization layer. + + Args: + inputs: Input tensor of shape [batch, H, W, C] + training: Whether in training mode + + Returns: + Normalized output tensor of same shape + """ + return self.bn(inputs, training=training) diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/concat.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/concat.py.jinja2 new file mode 100644 index 0000000..e774313 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/concat.py.jinja2 @@ -0,0 +1,32 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + Concatenation Layer + + Concatenates multiple tensors along a specified axis. + Commonly used to merge feature maps from different paths in the network. + + Parameters: + - Concatenation axis: {{ axis }} + + Shape: + - Input: List of tensors with compatible shapes + - Output: Concatenated tensor along axis {{ axis }} + """ + + def __init__(self, axis={{ axis }}): + """Initialize the concatenation layer.""" + super({{ class_name }}, self).__init__() + self.concat = tf.keras.layers.Concatenate(axis=axis) + + def call(self, inputs, training=None): + """ + Forward pass through the concatenation layer. + + Args: + inputs: List of input tensors + training: Whether in training mode + + Returns: + Concatenated tensor along axis {{ axis }} + """ + return self.concat(inputs) diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/conv2d.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/conv2d.py.jinja2 new file mode 100644 index 0000000..81366a9 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/conv2d.py.jinja2 @@ -0,0 +1,40 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + 2D Convolutional Layer + + Applies a 2D convolution over an input signal composed of several input channels. + + Parameters: + - Input channels: {{ in_channels }} + - Output channels: {{ out_channels }} + - Kernel size: {{ kernel_size }}x{{ kernel_size }} + - Strides: {{ strides }} + - Padding: {{ padding }} + + Shape: + - Input: [batch_size, H, W, {{ in_channels }}] (NHWC format) + - Output: [batch_size, H_out, W_out, {{ out_channels }}] + """ + + def __init__(self, filters={{ out_channels }}, kernel_size={{ kernel_size }}, strides={{ strides }}, padding='{{ padding }}'): + """Initialize the convolutional layer.""" + super({{ class_name }}, self).__init__() + self.conv = tf.keras.layers.Conv2D( + filters=filters, + kernel_size=kernel_size, + strides=strides, + padding=padding + ) + + def call(self, inputs, training=None): + """ + Forward pass through the convolutional layer. + + Args: + inputs: Input tensor of shape [batch, H, W, {{ in_channels }}] + training: Whether in training mode + + Returns: + Output tensor of shape [batch, H_out, W_out, {{ out_channels }}] + """ + return self.conv(inputs, training=training) diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/dropout.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/dropout.py.jinja2 new file mode 100644 index 0000000..4b32e1c --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/dropout.py.jinja2 @@ -0,0 +1,32 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + Dropout Regularization Layer + + Randomly sets input units to 0 with a frequency of rate during training. + Helps prevent overfitting. + + Parameters: + - Dropout rate: {{ rate }} + + Shape: + - Input: [batch_size, *] (any shape) + - Output: [batch_size, *] (same shape as input) + """ + + def __init__(self, rate={{ rate }}): + """Initialize the dropout layer.""" + super({{ class_name }}, self).__init__() + self.dropout = tf.keras.layers.Dropout(rate=rate) + + def call(self, inputs, training=None): + """ + Forward pass through the dropout layer. + + Args: + inputs: Input tensor + training: Whether in training mode + + Returns: + Output tensor with dropout applied during training + """ + return self.dropout(inputs, training=training) diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/flatten.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/flatten.py.jinja2 new file mode 100644 index 0000000..402a852 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/flatten.py.jinja2 @@ -0,0 +1,29 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + Flatten Layer + + Flattens the input tensor to 1D (preserving batch dimension). + Commonly used to transition from convolutional layers to fully connected layers. + + Shape: + - Input: [batch_size, H, W, C] + - Output: [batch_size, H*W*C]{% if out_features %} = [batch_size, {{ out_features }}]{% endif %} + """ + + def __init__(self): + """Initialize the flatten layer.""" + super({{ class_name }}, self).__init__() + self.flatten = tf.keras.layers.Flatten() + + def call(self, inputs, training=None): + """ + Forward pass through the flatten layer. + + Args: + inputs: Input tensor of shape [batch, H, W, C] + training: Whether in training mode + + Returns: + Output tensor of shape [batch, H*W*C] + """ + return self.flatten(inputs, training=training) diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/group_block.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/group_block.py.jinja2 new file mode 100644 index 0000000..b50c716 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/group_block.py.jinja2 @@ -0,0 +1,69 @@ +class {{ class_name }}(tf.keras.Model): + """ + {{ group_name }}{% if description %}: {{ description }}{% endif %} + + Custom group block with internal layers. +{% if has_multi_input %} + + Inputs: +{% for port in input_ports %} + - {{ port.externalPortLabel }}: {{ port.semantic }} (port: {{ port.externalPortId }}) +{% endfor %} +{% endif %} +{% if has_multi_output %} + + Outputs: +{% for port in output_ports %} + - {{ port.externalPortLabel }}: {{ port.semantic }} (port: {{ port.externalPortId }}) +{% endfor %} +{% endif %} + """ + + def __init__(self{% for param in init_params %}, {{ param.name }}={{ param.default }}{% endfor %}): + """Initialize {{ group_name }} with internal layers.""" + super({{ class_name }}, self).__init__() + + # Internal layers +{% for spec in internal_specs %} + self.{{ spec.layer_variable_name }} = {{ spec.class_name }}( +{% for key, value in spec.init_params.items() %} + {{ key }}={{ value }}{{ "," if not loop.last else "" }} +{% endfor %} + ) +{% endfor %} + + def call(self, inputs, training=None): + """ + Forward pass through {{ group_name }}. + + Args: +{% if has_multi_input %} + inputs: List or tuple of {{ num_inputs }} input tensors +{% else %} + inputs: Input tensor +{% endif %} + training: Whether in training mode + + Returns: +{% if has_multi_output %} + Tuple of {{ num_outputs }} output tensors: ({{ output_vars | join(', ') }}) +{% else %} + Output tensor +{% endif %} + """ +{% if has_multi_input %} + # Unpack multiple inputs +{% for i in range(num_inputs) %} + input_{{ i }} = inputs[{{ i }}] +{% endfor %} + +{% endif %} +{% for line in call_lines %} + {{ line }} +{% endfor %} + +{% if has_multi_output %} + return ({{ output_vars | join(', ') }}) +{% else %} + return {{ output_vars[0] if output_vars else 'inputs' }} +{% endif %} diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/linear.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/linear.py.jinja2 new file mode 100644 index 0000000..d23af28 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/linear.py.jinja2 @@ -0,0 +1,33 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + Dense/Fully Connected Layer + + Applies a linear transformation to the incoming data. + + Parameters: + - Input features: {{ in_features }} + - Output features: {{ out_features }} + - Use bias: {{ use_bias }} + + Shape: + - Input: [batch_size, {{ in_features }}] + - Output: [batch_size, {{ out_features }}] + """ + + def __init__(self, units={{ out_features }}, use_bias={{ use_bias }}): + """Initialize the dense layer.""" + super({{ class_name }}, self).__init__() + self.dense = tf.keras.layers.Dense(units=units, use_bias=use_bias) + + def call(self, inputs, training=None): + """ + Forward pass through the dense layer. + + Args: + inputs: Input tensor of shape [batch, {{ in_features }}] + training: Whether in training mode + + Returns: + Output tensor of shape [batch, {{ out_features }}] + """ + return self.dense(inputs, training=training) diff --git a/project/block_manager/services/nodes/templates/tensorflow/layers/maxpool.py.jinja2 b/project/block_manager/services/nodes/templates/tensorflow/layers/maxpool.py.jinja2 new file mode 100644 index 0000000..a7a0ec8 --- /dev/null +++ b/project/block_manager/services/nodes/templates/tensorflow/layers/maxpool.py.jinja2 @@ -0,0 +1,38 @@ +class {{ class_name }}(tf.keras.layers.Layer): + """ + 2D Max Pooling Layer + + Applies a 2D max pooling over an input signal. + Reduces spatial dimensions while preserving channel count. + + Parameters: + - Pool size: {{ pool_size }}x{{ pool_size }} + - Strides: {{ strides }} + - Padding: {{ padding }} + + Shape: + - Input: [batch_size, H, W, C] (NHWC format) + - Output: [batch_size, H/stride, W/stride, C] + """ + + def __init__(self, pool_size={{ pool_size }}, strides={{ strides }}, padding='{{ padding }}'): + """Initialize the max pooling layer.""" + super({{ class_name }}, self).__init__() + self.pool = tf.keras.layers.MaxPooling2D( + pool_size=pool_size, + strides=strides, + padding=padding + ) + + def call(self, inputs, training=None): + """ + Forward pass through the pooling layer. + + Args: + inputs: Input tensor of shape [batch, H, W, C] + training: Whether in training mode + + Returns: + Output tensor with reduced spatial dimensions + """ + return self.pool(inputs, training=training) diff --git a/project/block_manager/services/nodes/tensorflow/add.py b/project/block_manager/services/nodes/tensorflow/add.py index 5125d44..a8a56cb 100644 --- a/project/block_manager/services/nodes/tensorflow/add.py +++ b/project/block_manager/services/nodes/tensorflow/add.py @@ -1,7 +1,7 @@ """TensorFlow Add Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class AddNode(NodeDefinition): @@ -51,3 +51,27 @@ def validate_incoming_connection( def allows_multiple_inputs(self) -> bool: """Add nodes accept multiple input connections""" return True + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Add layer""" + sanitized_id = node_id.replace('-', '_') + class_name = 'AddBlock' + layer_var = f'{sanitized_id}_AddBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='add', + node_id=node_id, + init_params={}, + config_params=config, + input_shape_info={}, + output_shape_info={}, + template_context={} + ) + diff --git a/project/block_manager/services/nodes/tensorflow/attention.py b/project/block_manager/services/nodes/tensorflow/attention.py new file mode 100644 index 0000000..ee8fd69 --- /dev/null +++ b/project/block_manager/services/nodes/tensorflow/attention.py @@ -0,0 +1,108 @@ +"""TensorFlow Multi-Head Attention Node Definition""" + +from typing import Dict, List, Optional, Any +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec + + +class AttentionNode(NodeDefinition): + """Multi-Head Attention layer""" + + @property + def metadata(self) -> NodeMetadata: + return NodeMetadata( + type="attention", + label="Multi-Head Attention", + category="advanced", + color="var(--color-accent)", + icon="Zap", + description="Multi-head self-attention mechanism", + framework=Framework.TENSORFLOW + ) + + @property + def config_schema(self) -> List[ConfigField]: + return [ + ConfigField( + name="num_heads", + label="Number of Heads", + type="number", + default=8, + min=1, + description="Number of attention heads" + ), + ConfigField( + name="key_dim", + label="Key Dimension", + type="number", + default=64, + min=1, + description="Size of each attention head for query and key" + ), + ConfigField( + name="dropout", + label="Dropout", + type="number", + default=0.0, + min=0.0, + max=1.0, + description="Dropout probability" + ) + ] + + def compute_output_shape( + self, + input_shape: Optional[TensorShape], + config: Dict[str, Any] + ) -> Optional[TensorShape]: + # Multi-head attention preserves shape + if input_shape: + return TensorShape( + dims=input_shape.dims, + description="Attention output" + ) + return None + + def validate_incoming_connection( + self, + source_node_type: str, + source_output_shape: Optional[TensorShape], + target_config: Dict[str, Any] + ) -> Optional[str]: + # Attention accepts any input shape + return None + + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Multi-Head Attention layer""" + num_heads = config.get('num_heads', 8) + key_dim = config.get('key_dim', 64) + dropout = config.get('dropout', 0.0) + + sanitized_id = node_id.replace('-', '_') + class_name = 'MultiHeadAttentionLayer' + layer_var = f'{sanitized_id}_MultiHeadAttentionLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='attention', + node_id=node_id, + init_params={ + 'num_heads': num_heads, + 'key_dim': key_dim, + 'dropout': dropout + }, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={ + 'num_heads': num_heads, + 'key_dim': key_dim, + 'dropout': dropout + } + ) diff --git a/project/block_manager/services/nodes/tensorflow/batchnorm2d.py b/project/block_manager/services/nodes/tensorflow/batchnorm2d.py index 5fec405..5738056 100644 --- a/project/block_manager/services/nodes/tensorflow/batchnorm2d.py +++ b/project/block_manager/services/nodes/tensorflow/batchnorm2d.py @@ -1,7 +1,7 @@ """TensorFlow BatchNormalization Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class BatchNorm2DNode(NodeDefinition): @@ -70,3 +70,30 @@ def validate_incoming_connection( return "BatchNormalization requires input with at least 2 dimensions" return None + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for BatchNormalization layer""" + epsilon = config.get('epsilon', 1e-5) + momentum = config.get('momentum', 0.99) + + sanitized_id = node_id.replace('-', '_') + class_name = 'BatchNormBlock' + layer_var = f'{sanitized_id}_BatchNormBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='batchnorm', + node_id=node_id, + init_params={'epsilon': epsilon, 'momentum': momentum}, + config_params=config, + input_shape_info={}, + output_shape_info={}, + template_context={'epsilon': epsilon, 'momentum': momentum} + ) + diff --git a/project/block_manager/services/nodes/tensorflow/concat.py b/project/block_manager/services/nodes/tensorflow/concat.py index 4b545e7..5314ef5 100644 --- a/project/block_manager/services/nodes/tensorflow/concat.py +++ b/project/block_manager/services/nodes/tensorflow/concat.py @@ -1,7 +1,7 @@ """TensorFlow Concatenate Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class ConcatNode(NodeDefinition): @@ -62,3 +62,29 @@ def validate_incoming_connection( def allows_multiple_inputs(self) -> bool: """Concat nodes accept multiple input connections""" return True + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Concatenate layer""" + axis = config.get('axis', -1) + + sanitized_id = node_id.replace('-', '_') + class_name = 'ConcatBlock' + layer_var = f'{sanitized_id}_ConcatBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='concat', + node_id=node_id, + init_params={'axis': axis}, + config_params=config, + input_shape_info={}, + output_shape_info={}, + template_context={'axis': axis} + ) + diff --git a/project/block_manager/services/nodes/tensorflow/conv2d.py b/project/block_manager/services/nodes/tensorflow/conv2d.py index 976f7f6..d395fc9 100644 --- a/project/block_manager/services/nodes/tensorflow/conv2d.py +++ b/project/block_manager/services/nodes/tensorflow/conv2d.py @@ -1,7 +1,7 @@ """TensorFlow Conv2D Layer Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class Conv2DNode(NodeDefinition): @@ -122,3 +122,46 @@ def validate_incoming_connection( 4, "[batch, height, width, channels]" ) + + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Conv2D layer""" + out_channels = config.get('filters', 64) + kernel_size = config.get('kernel_size', 3) + strides = config.get('strides', 1) + padding = config.get('padding', 'valid') + + # TensorFlow uses NHWC format + in_channels = input_shape.dims[3] if input_shape and len(input_shape.dims) >= 4 else 3 + + sanitized_id = node_id.replace('-', '_') + class_name = 'Conv2DBlock' + layer_var = f'{sanitized_id}_Conv2DBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='conv2d', + node_id=node_id, + init_params={ + 'filters': out_channels, + 'kernel_size': kernel_size, + 'strides': strides, + 'padding': padding + }, + config_params=config, + input_shape_info={'in_channels': in_channels}, + output_shape_info={'out_channels': out_channels}, + template_context={ + 'in_channels': in_channels, + 'out_channels': out_channels, + 'kernel_size': kernel_size, + 'strides': strides, + 'padding': padding + } + ) diff --git a/project/block_manager/services/nodes/tensorflow/dropout.py b/project/block_manager/services/nodes/tensorflow/dropout.py index 6897bc8..467cb5f 100644 --- a/project/block_manager/services/nodes/tensorflow/dropout.py +++ b/project/block_manager/services/nodes/tensorflow/dropout.py @@ -1,7 +1,7 @@ """TensorFlow Dropout Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class DropoutNode(NodeDefinition): @@ -55,3 +55,29 @@ def validate_incoming_connection( ) -> Optional[str]: # Dropout accepts any input return None + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Dropout layer""" + rate = config.get('rate', 0.5) + + sanitized_id = node_id.replace('-', '_') + class_name = 'DropoutLayer' + layer_var = f'{sanitized_id}_DropoutLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='dropout', + node_id=node_id, + init_params={'rate': rate}, + config_params=config, + input_shape_info={}, + output_shape_info={}, + template_context={'rate': rate} + ) + diff --git a/project/block_manager/services/nodes/tensorflow/flatten.py b/project/block_manager/services/nodes/tensorflow/flatten.py index 3349ed0..9dba48d 100644 --- a/project/block_manager/services/nodes/tensorflow/flatten.py +++ b/project/block_manager/services/nodes/tensorflow/flatten.py @@ -1,7 +1,7 @@ """TensorFlow Flatten Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class FlattenNode(NodeDefinition): @@ -58,3 +58,27 @@ def validate_incoming_connection( return "Flatten requires input with at least 2 dimensions" return None + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Flatten layer""" + sanitized_id = node_id.replace('-', '_') + class_name = 'FlattenLayer' + layer_var = f'{sanitized_id}_FlattenLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='flatten', + node_id=node_id, + init_params={}, + config_params=config, + input_shape_info={}, + output_shape_info={}, + template_context={} + ) + diff --git a/project/block_manager/services/nodes/tensorflow/linear.py b/project/block_manager/services/nodes/tensorflow/linear.py index 81963f0..5546a96 100644 --- a/project/block_manager/services/nodes/tensorflow/linear.py +++ b/project/block_manager/services/nodes/tensorflow/linear.py @@ -1,7 +1,7 @@ """TensorFlow Dense (Linear) Layer Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class LinearNode(NodeDefinition): @@ -92,3 +92,35 @@ def validate_incoming_connection( return "Dense layer requires input with at least 2 dimensions [batch, features, ...]" return None + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Dense/Linear layer""" + out_features = config.get('units', 128) + use_bias = config.get('use_bias', True) + in_features = input_shape.dims[1] if input_shape and len(input_shape.dims) >= 2 else 512 + + sanitized_id = node_id.replace('-', '_') + class_name = 'DenseLayer' + layer_var = f'{sanitized_id}_DenseLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='linear', + node_id=node_id, + init_params={'units': out_features, 'use_bias': use_bias}, + config_params=config, + input_shape_info={'in_features': in_features}, + output_shape_info={'out_features': out_features}, + template_context={ + 'in_features': in_features, + 'out_features': out_features, + 'use_bias': use_bias + } + ) + diff --git a/project/block_manager/services/nodes/tensorflow/maxpool2d.py b/project/block_manager/services/nodes/tensorflow/maxpool2d.py index d831900..e32cfa8 100644 --- a/project/block_manager/services/nodes/tensorflow/maxpool2d.py +++ b/project/block_manager/services/nodes/tensorflow/maxpool2d.py @@ -1,7 +1,7 @@ """TensorFlow MaxPooling2D Node Definition""" from typing import Dict, List, Optional, Any -from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec class MaxPool2DNode(NodeDefinition): @@ -95,3 +95,31 @@ def validate_incoming_connection( 4, "[batch, height, width, channels]" ) + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for MaxPooling2D layer""" + pool_size = config.get('pool_size', 2) + strides = config.get('strides', 2) + padding = config.get('padding', 'valid') + + sanitized_id = node_id.replace('-', '_') + class_name = 'MaxPoolBlock' + layer_var = f'{sanitized_id}_MaxPoolBlock' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='maxpool', + node_id=node_id, + init_params={'pool_size': pool_size, 'strides': strides, 'padding': padding}, + config_params=config, + input_shape_info={}, + output_shape_info={}, + template_context={'pool_size': pool_size, 'strides': strides, 'padding': padding} + ) + diff --git a/project/block_manager/services/nodes/tensorflow/relu.py b/project/block_manager/services/nodes/tensorflow/relu.py new file mode 100644 index 0000000..231dd19 --- /dev/null +++ b/project/block_manager/services/nodes/tensorflow/relu.py @@ -0,0 +1,70 @@ +"""TensorFlow ReLU Activation Node Definition""" + +from typing import Dict, List, Optional, Any +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec + + +class ReLUNode(NodeDefinition): + """ReLU activation function layer""" + + @property + def metadata(self) -> NodeMetadata: + return NodeMetadata( + type="relu", + label="ReLU", + category="basic", + color="var(--color-primary)", + icon="Zap", + description="ReLU activation function", + framework=Framework.TENSORFLOW + ) + + @property + def config_schema(self) -> List[ConfigField]: + return [] # ReLU has no configurable parameters in TensorFlow + + def compute_output_shape( + self, + input_shape: Optional[TensorShape], + config: Dict[str, Any] + ) -> Optional[TensorShape]: + # ReLU preserves shape + if input_shape: + return TensorShape( + dims=input_shape.dims, + description="ReLU activated" + ) + return None + + def validate_incoming_connection( + self, + source_node_type: str, + source_output_shape: Optional[TensorShape], + target_config: Dict[str, Any] + ) -> Optional[str]: + # ReLU accepts any input shape + return None + + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for ReLU layer""" + sanitized_id = node_id.replace('-', '_') + class_name = 'ReLULayer' + layer_var = f'{sanitized_id}_ReLULayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='relu', + node_id=node_id, + init_params={}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={} + ) diff --git a/project/block_manager/services/nodes/tensorflow/softmax.py b/project/block_manager/services/nodes/tensorflow/softmax.py new file mode 100644 index 0000000..6d35284 --- /dev/null +++ b/project/block_manager/services/nodes/tensorflow/softmax.py @@ -0,0 +1,80 @@ +"""TensorFlow Softmax Activation Node Definition""" + +from typing import Dict, List, Optional, Any +from ..base import NodeDefinition, NodeMetadata, ConfigField, TensorShape, Framework, LayerCodeSpec + + +class SoftmaxNode(NodeDefinition): + """Softmax activation function layer""" + + @property + def metadata(self) -> NodeMetadata: + return NodeMetadata( + type="softmax", + label="Softmax", + category="basic", + color="var(--color-primary)", + icon="Activity", + description="Softmax activation function", + framework=Framework.TENSORFLOW + ) + + @property + def config_schema(self) -> List[ConfigField]: + return [ + ConfigField( + name="axis", + label="Axis", + type="number", + default=-1, + description="Axis along which softmax will be computed" + ) + ] + + def compute_output_shape( + self, + input_shape: Optional[TensorShape], + config: Dict[str, Any] + ) -> Optional[TensorShape]: + # Softmax preserves shape + if input_shape: + return TensorShape( + dims=input_shape.dims, + description="Softmax probabilities" + ) + return None + + def validate_incoming_connection( + self, + source_node_type: str, + source_output_shape: Optional[TensorShape], + target_config: Dict[str, Any] + ) -> Optional[str]: + # Softmax accepts any input shape + return None + + def get_tensorflow_code_spec( + self, + node_id: str, + config: Dict[str, Any], + input_shape: Optional[TensorShape], + output_shape: Optional[TensorShape] + ) -> LayerCodeSpec: + """Generate TensorFlow code specification for Softmax layer""" + axis = config.get('axis', -1) + + sanitized_id = node_id.replace('-', '_') + class_name = 'SoftmaxLayer' + layer_var = f'{sanitized_id}_SoftmaxLayer' + + return LayerCodeSpec( + class_name=class_name, + layer_variable_name=layer_var, + node_type='softmax', + node_id=node_id, + init_params={'axis': axis}, + config_params=config, + input_shape_info={'dims': input_shape.dims if input_shape else []}, + output_shape_info={'dims': output_shape.dims if output_shape else []}, + template_context={'axis': axis} + ) diff --git a/project/block_manager/services/tensorflow_codegen.py b/project/block_manager/services/tensorflow_codegen.py index 23dac16..7c7c5bd 100644 --- a/project/block_manager/services/tensorflow_codegen.py +++ b/project/block_manager/services/tensorflow_codegen.py @@ -7,620 +7,24 @@ from collections import deque import logging -# Import shared exceptions from PyTorch codegen (framework-agnostic) -from .enhanced_pytorch_codegen import ( - GroupDefinitionNotFoundError, - ShapeMismatchError, - CyclicDependencyError, - UnsupportedNodeTypeError, - ShapeInferenceError, - MissingShapeDataError -) +# New template-based code generation +from .codegen.tensorflow_orchestrator import TensorFlowCodeOrchestrator + +# NOTE: Legacy imports removed - all code generation now delegated to TensorFlowCodeOrchestrator +# The classes below were only used in legacy code that no longer executes: +# - GroupBlockShapeComputer, GroupDefinitionNotFoundError, ShapeMismatchError +# - CyclicDependencyError, UnsupportedNodeTypeError, ShapeInferenceError +# - MissingShapeDataError, safe_get_shape_data # Configure logging logger = logging.getLogger(__name__) -# ============================================ -# Helper Functions for Shape Inference -# ============================================ - -def safe_get_shape_data( - shape_map: Dict[str, Dict[str, Any]], - node_id: str, - upstream_node_id: str, - required_keys: List[str], - default_values: Optional[Dict[str, Any]] = None -) -> Dict[str, Any]: - """ - Safely retrieve shape data from upstream node with proper error handling. - - Args: - shape_map: Dictionary mapping node IDs to shape information - node_id: ID of the current node requesting shape data - upstream_node_id: ID of the upstream node to get shape from - required_keys: List of required keys in the shape data - default_values: Optional default values to use if keys are missing - - Returns: - Dictionary with shape data - - Raises: - MissingShapeDataError: If required keys are missing and no defaults provided - """ - if upstream_node_id not in shape_map: - if default_values: - return default_values - raise MissingShapeDataError( - node_id=node_id, - upstream_node_id=upstream_node_id, - missing_keys=required_keys - ) - - upstream_shape = shape_map[upstream_node_id] - missing_keys = [key for key in required_keys if key not in upstream_shape] - - if missing_keys: - if default_values: - # Use defaults for missing keys but keep existing values - result = upstream_shape.copy() - for key in missing_keys: - if key in default_values: - result[key] = default_values[key] - return result - raise MissingShapeDataError( - node_id=node_id, - upstream_node_id=upstream_node_id, - missing_keys=missing_keys - ) - - return upstream_shape - - -class GroupBlockShapeComputer: - """ - Computes shapes for group blocks by analyzing their internal structure. - TensorFlow version using NHWC format (batch, height, width, channels). - """ - - def __init__(self, group_definitions: Dict[str, Any]): - """ - Initialize the shape computer. - - Args: - group_definitions: Dictionary mapping definition IDs to group definitions - """ - self.group_definitions = group_definitions - - def compute_output_shape( - self, - group_def_id: str, - input_shape: Dict[str, Any] - ) -> Tuple[Dict[str, Any], List[Exception]]: - """ - Compute the output shape of a group block given its input shape. - - Args: - group_def_id: ID of the group definition - input_shape: Input shape dictionary - - Returns: - Tuple of (output shape dictionary, list of errors) - """ - if group_def_id not in self.group_definitions: - error = GroupDefinitionNotFoundError( - node_id=f"group_block_{group_def_id}", - definition_id=group_def_id - ) - return {}, [error] - - definition = self.group_definitions[group_def_id] - internal_structure = definition.get('internal_structure', {}) - internal_nodes = internal_structure.get('nodes', []) - internal_edges = internal_structure.get('edges', []) - port_mappings = internal_structure.get('portMappings', []) - - # Compute internal shapes - internal_shape_map, errors = self.compute_internal_shapes( - internal_nodes, - internal_edges, - port_mappings, - input_shape, - definition.get('name', 'UnnamedBlock') - ) - - # Find output port and return its shape - output_ports = [pm for pm in port_mappings if pm['type'] == 'output'] - if output_ports: - output_node_id = output_ports[0]['internalNodeId'] - if output_node_id in internal_shape_map: - return internal_shape_map[output_node_id], errors - - # Fallback: return input shape - return input_shape.copy(), errors - - def compute_internal_shapes( - self, - internal_nodes: List[Dict[str, Any]], - internal_edges: List[Dict[str, Any]], - port_mappings: List[Dict[str, Any]], - input_shape: Dict[str, Any], - block_name: str - ) -> Tuple[Dict[str, Dict[str, Any]], List[Exception]]: - """ - Compute shapes for all internal nodes in a group block. - - Args: - internal_nodes: List of internal node definitions - internal_edges: List of internal edges - port_mappings: Port mappings (input/output) - input_shape: Input shape for the block - block_name: Name of the block for error messages - - Returns: - Tuple of (shape map for internal nodes, list of errors) - """ - # Sort nodes topologically - sorted_nodes = topological_sort(internal_nodes, internal_edges) - - # Initialize shape map with input port shapes - shape_map = {} - input_ports = [pm for pm in port_mappings if pm['type'] == 'input'] - for port in input_ports: - internal_node_id = port['internalNodeId'] - shape_map[internal_node_id] = input_shape.copy() - - # Use the main infer_shapes function but pass None for group_definitions - # to avoid recursive group block resolution - internal_shape_map, errors = infer_shapes(sorted_nodes, internal_edges, None) - - # Merge input port shapes - for node_id, shape in shape_map.items(): - if node_id in internal_shape_map: - internal_shape_map[node_id].update(shape) - else: - internal_shape_map[node_id] = shape - - return internal_shape_map, errors - - -class TensorFlowBlockGenerator: - """ - Generator for TensorFlow/Keras tf.keras.Model code for group blocks. - - Converts GroupBlockDefinition into reusable tf.keras.Model subclasses - with proper initialization and call method logic. - """ - - def __init__( - self, - group_definitions: List[Dict[str, Any]], - shape_computer: Optional[GroupBlockShapeComputer] = None - ): - """ - Initialize the block generator. - - Args: - group_definitions: List of GroupBlockDefinition dictionaries - shape_computer: Optional shape computer for internal shape inference - """ - self.group_definitions = {defn['id']: defn for defn in group_definitions} - self.generated_classes = {} # Cache generated class code - self.shape_computer = shape_computer or GroupBlockShapeComputer(self.group_definitions) - - def generate_all_block_classes(self) -> str: - """ - Generate all block class definitions. - - Returns: - String containing all block class definitions - """ - if not self.group_definitions: - return "" - - code_parts = [] - code_parts.append("# ============================================") - code_parts.append("# Custom Block Definitions") - code_parts.append("# ============================================\n") - - for defn_id, definition in self.group_definitions.items(): - block_class = self.generate_block_class(definition) - code_parts.append(block_class) - code_parts.append("\n") - - return "\n".join(code_parts) - - def generate_block_class( - self, - definition: Dict[str, Any], - example_input_shape: Optional[Dict[str, Any]] = None - ) -> str: - """ - Generate tf.keras.Model subclass for a single block definition. - - Args: - definition: GroupBlockDefinition dictionary - example_input_shape: Optional example input shape for computing internal shapes - - Returns: - String containing the complete block class definition - """ - block_name = definition['name'] - class_name = self._to_class_name(block_name) - description = definition.get('description', '') - - # Get internal structure - internal_structure = definition.get('internal_structure', {}) - internal_nodes = internal_structure.get('nodes', []) - internal_edges = internal_structure.get('edges', []) - port_mappings = internal_structure.get('portMappings', []) - - # Sort internal nodes topologically - sorted_nodes = topological_sort(internal_nodes, internal_edges) - - # Compute internal shapes if example provided - internal_shape_map = {} - if example_input_shape: - internal_shape_map, _ = self.shape_computer.compute_internal_shapes( - internal_nodes, - internal_edges, - port_mappings, - example_input_shape, - block_name - ) - else: - # Fallback to old behavior without shape computer - internal_shape_map, _ = infer_shapes(sorted_nodes, internal_edges) - - # Generate __init__ method - init_method = self._generate_init_method(sorted_nodes, internal_shape_map, port_mappings) - - # Generate call method - call_method = self._generate_call_method( - sorted_nodes, internal_edges, internal_shape_map, port_mappings - ) - - # Build class docstring - docstring = self._generate_block_docstring( - block_name, description, port_mappings, sorted_nodes - ) - - # Assemble the complete class - class_code = f'''class {class_name}(keras.Model): - """{docstring}""" - -{init_method} - -{call_method}''' - - # Cache the generated class - self.generated_classes[definition['id']] = class_name - - return class_code - - def _generate_init_method( - self, - nodes: List[Dict[str, Any]], - shape_map: Dict[str, Dict[str, Any]], - port_mappings: List[Dict[str, Any]] - ) -> str: - """Generate __init__ method with layer instantiation.""" - lines = [] - - # Detect which shape parameters are needed by scanning nodes - needs_in_channels = False - needs_in_features = False - needs_num_features = False - - for node in nodes: - node_type = get_node_type(node) - if node_type in ('input', 'dataloader', 'output'): - continue - if node_type == 'conv2d': - needs_in_channels = True - elif node_type == 'linear': - needs_in_features = True - elif node_type in ('batchnorm', 'batchnorm2d'): - needs_num_features = True - - # Generate __init__ signature with detected parameters - params = [] - if needs_in_channels: - params.append("in_channels=None") - if needs_in_features: - params.append("in_features=None") - if needs_num_features: - params.append("num_features=None") - - if params: - lines.append(f" def __init__(self, {', '.join(params)}):") - else: - lines.append(" def __init__(self):") - - lines.append(' """Initialize all internal layers."""') - lines.append(f" super().__init__()") - lines.append("") - - # Track which nodes need to be instantiated and which is first of each type - layer_count = {} - first_layer_of_type = {} - - for idx, node in enumerate(nodes): - node_id = node['id'] - node_type = get_node_type(node) - config = node.get('data', {}).get('config', {}) - shape_info = shape_map.get(node_id, {}) - - # Skip input/output nodes - if node_type in ('input', 'dataloader', 'output'): - continue - - # Track if this is the first layer of its type - is_first = node_type not in first_layer_of_type - if is_first: - first_layer_of_type[node_type] = node_id - - # Generate layer instantiation - layer_name = self._get_internal_layer_name(node_type, node_id, layer_count) - layer_class_name = self._get_layer_class_name_for_node(node_type, config) - - # Generate instantiation with proper arguments - instantiation = self._generate_layer_instantiation_line( - layer_name, layer_class_name, node_type, shape_info, config, is_first - ) - - if instantiation: - lines.append(f" {instantiation}") - - return "\n".join(lines) - - def _generate_call_method( - self, - nodes: List[Dict[str, Any]], - edges: List[Dict[str, Any]], - shape_map: Dict[str, Dict[str, Any]], - port_mappings: List[Dict[str, Any]] - ) -> str: - """Generate call method with internal connection logic.""" - lines = [] - - # Determine input parameters from port mappings - input_ports = [pm for pm in port_mappings if pm['type'] == 'input'] - output_ports = [pm for pm in port_mappings if pm['type'] == 'output'] - - # Generate method signature - if len(input_ports) == 1: - lines.append(" def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor:") - else: - param_names = [f"input_{i}" for i in range(len(input_ports))] - params = ", ".join([f"{name}: tf.Tensor" for name in param_names]) - lines.append(f" def call(self, {params}, training: Optional[bool] = None) -> tf.Tensor:") - - lines.append(' """') - lines.append(' Forward pass through the block.') - lines.append('') - lines.append(' Args:') - if len(input_ports) == 1: - lines.append(' inputs: Input tensor in NHWC format') - else: - for i, port in enumerate(input_ports): - label = port.get('externalPortLabel', f'input_{i}') - lines.append(f' input_{i}: {label}') - lines.append(' training: Whether in training mode') - lines.append('') - lines.append(' Returns:') - if len(output_ports) == 1: - lines.append(' Output tensor') - else: - lines.append(' Tuple of output tensors') - lines.append(' """') - - # Build edge map for finding inputs - edge_map = {} - for edge in edges: - target = edge.get('target') - source = edge.get('source') - if target not in edge_map: - edge_map[target] = [] - edge_map[target].append(source) - - # Map internal node IDs to variable names - var_map = {} - layer_count = {} - - # Map input ports to initial variables - for i, port in enumerate(input_ports): - internal_node_id = port['internalNodeId'] - if len(input_ports) == 1: - var_map[internal_node_id] = 'inputs' - else: - var_map[internal_node_id] = f'input_{i}' - - # Generate forward pass for each internal node - for node in nodes: - node_id = node['id'] - node_type = get_node_type(node) - config = node.get('data', {}).get('config', {}) - - # Skip input/output nodes - if node_type in ('input', 'dataloader', 'output'): - # Input nodes are already mapped - if node_id not in var_map: - var_map[node_id] = 'inputs' - continue - - # Get layer name - layer_name = self._get_internal_layer_name(node_type, node_id, layer_count) - - # Get input variable(s) - incoming = edge_map.get(node_id, []) - if not incoming: - # No incoming edges, might be an input node we missed - input_var = 'inputs' - elif len(incoming) == 1: - input_var = var_map.get(incoming[0], 'inputs') - else: - # Multiple inputs (for concat, add, etc.) - input_vars = [var_map.get(src, 'inputs') for src in incoming] - input_var = f"[{', '.join(input_vars)}]" - - # Generate output variable name (sanitize node_id to avoid hyphens) - output_var = f"x_{node_id[:8].replace('-', '_')}" - var_map[node_id] = output_var - - # Generate forward line with training parameter for layers that need it - if node_type in ('dropout', 'batchnorm', 'batchnorm2d'): - lines.append(f" {output_var} = self.{layer_name}({input_var}, training=training)") - else: - lines.append(f" {output_var} = self.{layer_name}({input_var})") - - # Map output ports to return values - if len(output_ports) == 1: - output_node_id = output_ports[0]['internalNodeId'] - output_var = var_map.get(output_node_id, 'inputs') - lines.append(f" return {output_var}") - else: - output_vars = [] - for port in output_ports: - output_node_id = port['internalNodeId'] - output_vars.append(var_map.get(output_node_id, 'inputs')) - lines.append(f" return ({', '.join(output_vars)})") - - return "\n".join(lines) - - def _generate_block_docstring( - self, - block_name: str, - description: str, - port_mappings: List[Dict[str, Any]], - nodes: List[Dict[str, Any]] - ) -> str: - """Generate comprehensive docstring for block class.""" - lines = [] - lines.append(f"Custom Block: {block_name}") - lines.append("") - - if description: - lines.append(description) - lines.append("") - - lines.append("This block encapsulates a reusable subgraph of layers.") - lines.append("") - lines.append("Note: TensorFlow uses NHWC format (batch, height, width, channels)") - lines.append("") - - # Document ports - input_ports = [pm for pm in port_mappings if pm['type'] == 'input'] - output_ports = [pm for pm in port_mappings if pm['type'] == 'output'] - - if input_ports: - lines.append("Input Ports:") - for port in input_ports: - label = port.get('externalPortLabel', 'input') - lines.append(f" - {label}") - - if output_ports: - lines.append("") - lines.append("Output Ports:") - for port in output_ports: - label = port.get('externalPortLabel', 'output') - lines.append(f" - {label}") - - lines.append("") - lines.append(f"Internal Layers: {len([n for n in nodes if get_node_type(n) not in ('input', 'dataloader', 'output')])}") - - return "\n ".join(lines) - - def _generate_layer_instantiation_line( - self, - layer_name: str, - layer_class_name: str, - node_type: str, - shape_info: Dict[str, Any], - config: Dict[str, Any], - is_first: bool = False - ) -> str: - """ - Generate layer instantiation line for TensorFlow/Keras layers. - - TensorFlow/Keras layer classes have all configuration baked into their - class definitions, so __init__ methods take no parameters. This differs - from PyTorch where layers need input dimensions in the constructor. - """ - # TensorFlow layers don't need input shape parameters in constructor - # All configuration is already baked into the layer class definition - # Just instantiate with no arguments - return f"self.{layer_name} = {layer_class_name}()" - - def _get_internal_layer_name( - self, - node_type: str, - node_id: str, - layer_count: Dict[str, int] - ) -> str: - """Generate unique layer variable name for internal node.""" - # Use node_id suffix for uniqueness (sanitize to avoid hyphens) - suffix = node_id[:8].replace('-', '_') - base_name = node_type.replace('_', '') - - # Track count for this type - if node_type not in layer_count: - layer_count[node_type] = 0 - layer_count[node_type] += 1 - - return f"{base_name}_{suffix}" - - def _get_layer_class_name_for_node( - self, - node_type: str, - config: Dict[str, Any] - ) -> str: - """Get the layer class name that will be used in the main model.""" - # These should match the class names generated by generate_layer_class - type_name = node_type.replace('_', '').replace('2d', '2D').replace('3d', '3D').title() - - if node_type == 'conv2d': - filters = config.get('filters', 64) - kernel = config.get('kernel_size', 3) - return f"{type_name}Layer_{filters}filters_{kernel}x{kernel}" - elif node_type == 'linear': - units = config.get('units', 128) - return f"DenseLayer_{units}units" - elif node_type in ('maxpool2d', 'maxpool'): - pool_size = config.get('pool_size', 2) - return f"MaxPool2DLayer_{pool_size}x{pool_size}" - elif node_type == 'custom': - name = config.get('name', 'CustomLayer') - safe_name = name.replace(' ', '_').replace('-', '_') - return f"CustomLayer_{safe_name}" - else: - # For other types, we'll need to generate a generic name - # This will be handled by the main code generation - return f"{type_name}Layer" - - def _to_class_name(self, name: str) -> str: - """Convert block name to valid Python class name.""" - import re - # Remove special characters and convert to PascalCase - name = re.sub(r'[^a-zA-Z0-9]', ' ', name) - name = ''.join(word.capitalize() for word in name.split()) - if not name: - return 'CustomBlock' - if name[0].isdigit(): - name = 'Block' + name - return name + 'Block' - - def get_block_class_name(self, definition_id: str) -> Optional[str]: - """ - Get the generated class name for a block definition. - - Args: - definition_id: ID of the GroupBlockDefinition - - Returns: - Class name if generated, None otherwise - """ - return self.generated_classes.get(definition_id) +# ==================== LEGACY CLASS REMOVED ==================== +# TensorFlowBlockGenerator class has been removed and replaced with: +# - TensorFlowGroupBlockGenerator in codegen/tensorflow_group_generator.py (for group blocks) +# - TensorFlowCodeOrchestrator in codegen/tensorflow_orchestrator.py (for overall code generation) +# =============================================================== def generate_tensorflow_code( @@ -642,1969 +46,13 @@ def generate_tensorflow_code( Returns: Tuple of (dictionary with keys: 'model', 'train', 'dataset', 'config', list of errors) """ - # Topologically sort nodes - sorted_nodes = topological_sort(nodes, edges) - - # Convert group_definitions list to dict for shape inference - group_defs_dict = None - if group_definitions: - group_defs_dict = {defn['id']: defn for defn in group_definitions} - - # Infer shapes through the graph with group block support - shape_map, shape_errors = infer_shapes(sorted_nodes, edges, group_defs_dict) - - # Validate computed shapes for critical issues - validation_errors = validate_shape_map(sorted_nodes, shape_map) - if validation_errors: - logger.warning(f"Shape validation found {len(validation_errors)} potential issues") - shape_errors.extend(validation_errors) - - # Initialize block generator if we have group definitions - block_generator = None - if group_definitions: - # Create shape computer for block generator - shape_computer = GroupBlockShapeComputer(group_defs_dict) if group_defs_dict else None - block_generator = TensorFlowBlockGenerator(group_definitions, shape_computer) - - # Generate different components - model_code = generate_model_file(sorted_nodes, edges, project_name, shape_map, block_generator) - train_code = generate_training_script(project_name) - dataset_code = generate_dataset_class(nodes) - config_code = generate_config_file(nodes) - - # Return generated code with any shape inference errors - return { - 'model': model_code, - 'train': train_code, - 'dataset': dataset_code, - 'config': config_code - }, shape_errors - - -def generate_single_layer_class( - node: Dict[str, Any], - node_index: int = 0, - shape_info: Optional[Dict[str, Any]] = None -) -> str: - """ - Generate professional class-based code for a single layer. - Used for individual node preview in the visual editor. - - Args: - node: Node dictionary with type, data, config - node_index: Index for layer naming (default: 0) - shape_info: Optional shape information dict. If None, extracted from node. - - Returns: - String containing the complete layer class definition - """ - # Extract node information - node_type = get_node_type(node) - config = node.get('data', {}).get('config', {}) - - # Extract or infer shape information - if shape_info is None: - shape_info = extract_shape_info_from_node(node) - - # Skip nodes that don't generate layers - if node_type in ('input', 'dataloader', 'output'): - return f'''# {node_type.upper()} Node -# This is handled automatically during model execution -# Input shape (NHWC): {shape_info.get('out_channels', '?')} channels or {shape_info.get('out_units', '?')} units''' - - # Generate the layer class using existing function - layer_class = generate_layer_class(node, node_index, config, node_type, shape_info) - - if layer_class: - return layer_class - else: - return f'''# Unsupported layer type: {node_type} -# Please use the full export to generate complete model code''' - - -def extract_shape_info_from_node(node: Dict[str, Any]) -> Dict[str, Any]: - """ - Extract shape information from a single node's metadata. - TensorFlow uses NHWC format (batch, height, width, channels). - - Args: - node: Node dictionary - - Returns: - Dictionary with shape information (in_channels, out_channels, in_units, out_units, etc.) - """ - shape_info = {} - node_type = get_node_type(node) - config = node.get('data', {}).get('config', {}) - - # Try to get shape from node metadata - input_shape = node.get('data', {}).get('inputShape', {}) - output_shape = node.get('data', {}).get('outputShape', {}) - - # Extract from inputShape/outputShape if available (NHWC format) - if input_shape and isinstance(input_shape, dict): - dims = input_shape.get('dims', []) - if len(dims) >= 4: # NHWC format - shape_info['in_height'] = dims[1] - shape_info['in_width'] = dims[2] - shape_info['in_channels'] = dims[3] - elif len(dims) >= 2: - shape_info['in_units'] = dims[1] - - if output_shape and isinstance(output_shape, dict): - dims = output_shape.get('dims', []) - if len(dims) >= 4: # NHWC format - shape_info['out_height'] = dims[1] - shape_info['out_width'] = dims[2] - shape_info['out_channels'] = dims[3] - elif len(dims) >= 2: - shape_info['out_units'] = dims[1] - - # Infer from config if not in metadata - if node_type == 'conv2d': - if 'in_channels' not in shape_info: - shape_info['in_channels'] = 3 # Default - if 'out_channels' not in shape_info: - shape_info['out_channels'] = config.get('filters', 64) - # Try to estimate output dimensions if not provided - if 'out_height' not in shape_info: - shape_info['out_height'] = '?' - if 'out_width' not in shape_info: - shape_info['out_width'] = '?' - - elif node_type == 'linear': - if 'in_units' not in shape_info: - shape_info['in_units'] = 512 # Default - if 'out_units' not in shape_info: - shape_info['out_units'] = config.get('units', 128) - - elif node_type in ('batchnorm', 'batchnorm2d'): - # BatchNorm preserves shape - if 'out_channels' not in shape_info: - shape_info['out_channels'] = shape_info.get('in_channels', 64) - - elif node_type == 'flatten': - if 'out_units' not in shape_info: - # Estimate based on typical conv output (NHWC) - height = shape_info.get('in_height', 7) - width = shape_info.get('in_width', 7) - channels = shape_info.get('in_channels', 512) - if isinstance(height, int) and isinstance(width, int) and isinstance(channels, int): - shape_info['out_units'] = height * width * channels - else: - shape_info['out_units'] = '?' - - return shape_info - - -def topological_sort(nodes: List[Dict], edges: List[Dict]) -> List[Dict]: - """Sort nodes in topological order based on edges using Kahn's algorithm""" - node_map = {node['id']: node for node in nodes} - - # Build adjacency list and in-degree count - graph = {node['id']: [] for node in nodes} - in_degree = {node['id']: 0 for node in nodes} - - for edge in edges: - source = edge.get('source') - target = edge.get('target') - if source in graph and target in graph: - graph[source].append(target) - in_degree[target] += 1 - - # Kahn's algorithm - queue = deque([node_id for node_id, degree in in_degree.items() if degree == 0]) - sorted_ids = [] - - while queue: - node_id = queue.popleft() - sorted_ids.append(node_id) - - for neighbor in graph[node_id]: - in_degree[neighbor] -= 1 - if in_degree[neighbor] == 0: - queue.append(neighbor) - - # Return nodes in sorted order - return [node_map[node_id] for node_id in sorted_ids if node_id in node_map] - - -def extract_output_shape_from_metadata(node: Dict[str, Any]) -> Optional[Dict[str, Any]]: - """ - Extract output shape from node's frontend-provided metadata (TensorFlow/NHWC version). - - The frontend computes output shapes accurately during the visual design phase - and stores them in node.data.outputShape. This function extracts those - pre-computed shapes, which are considered authoritative. - - Args: - node: Node dictionary with potential data.outputShape metadata - - Returns: - Dictionary with shape keys (out_channels, out_features, etc.) or None if - metadata is incomplete/missing - """ - output_shape = node.get('data', {}).get('outputShape', {}) - if not output_shape or not isinstance(output_shape, dict): - return None - - dims = output_shape.get('dims', []) - if not dims: - return None - - shape_info = {} - - # TensorFlow uses NHWC format: [batch, height, width, channels] - # Note: This is different from PyTorch's NCHW format! - if len(dims) == 4: - shape_info['out_height'] = dims[1] - shape_info['out_width'] = dims[2] - shape_info['out_channels'] = dims[3] - elif len(dims) == 2: # [batch, features] - for Dense/Flatten output - shape_info['out_features'] = dims[1] - else: - # Unusual shape format - log for debugging but don't fail - logger.debug(f"Unusual output shape dims: {dims}") - return None - - return shape_info - - -def infer_shapes( - nodes: List[Dict], - edges: List[Dict], - group_definitions: Optional[Dict[str, Any]] = None -) -> Tuple[Dict[str, Dict[str, Any]], List[Exception]]: - """ - Infer input/output shapes for each layer in the graph. - TensorFlow uses NHWC format (batch, height, width, channels). - Enhanced to handle group blocks properly. - - Args: - nodes: List of node dictionaries - edges: List of edge dictionaries - group_definitions: Optional map of group definition IDs to definitions - - Returns: - Tuple of (dictionary mapping node_id to shape info, list of errors) - """ - shape_map = {} - errors = [] - - # Initialize shape computer for group blocks - shape_computer = None - if group_definitions: - shape_computer = GroupBlockShapeComputer(group_definitions) - - # Build edge map for finding inputs - edge_map = {} - for edge in edges: - target = edge.get('target') - source = edge.get('source') - if target not in edge_map: - edge_map[target] = [] - edge_map[target].append(source) - - # Process nodes in order - for node in nodes: - node_id = node['id'] - node_type = get_node_type(node) - config = node.get('data', {}).get('config', {}) - - # Get incoming edges - incoming = edge_map.get(node_id, []) - - # ========== PHASE 1: Extract output metadata (if available) ========== - # Frontend provides accurate output shapes in metadata - metadata_shape = extract_output_shape_from_metadata(node) - shape_info = metadata_shape if metadata_shape else {} - - # ========== PHASE 2: Compute input dimensions from upstream nodes ========== - # Input dimensions ALWAYS come from upstream, regardless of metadata - # This is critical for layers like Conv2D, Dense, BatchNorm - - if node_type == 'input': - # Input nodes have no upstream - parse from config if metadata doesn't exist - if not metadata_shape: - shape_str = config.get('shape', '[1, 224, 224, 3]') - try: - import json - shape = json.loads(shape_str) - if len(shape) >= 4: - shape_info['out_height'] = shape[1] # NHWC format - shape_info['out_width'] = shape[2] - shape_info['out_channels'] = shape[3] - elif len(shape) >= 2: - shape_info['out_units'] = shape[1] - except (json.JSONDecodeError, ValueError, KeyError, IndexError, TypeError) as e: - logger.warning( - f"Failed to parse input shape for node {node_id}: {e}. " - f"Using default shape [1, 224, 224, 3] (NHWC)" - ) - errors.append(ShapeInferenceError( - node_id=node_id, - node_type=node_type, - reason=f"Failed to parse shape configuration: {str(e)}", - suggestion="Check that the input shape is a valid JSON array like [1, 224, 224, 3]" - )) - shape_info['out_height'] = 224 - shape_info['out_width'] = 224 - shape_info['out_channels'] = 3 - - elif node_type == 'conv2d': - # Get input channels from upstream layer (ALWAYS required) - if incoming and incoming[0] in shape_map: - try: - upstream_shape = safe_get_shape_data( - shape_map=shape_map, - node_id=node_id, - upstream_node_id=incoming[0], - required_keys=['out_channels'], - default_values={'out_channels': 3} - ) - shape_info['in_channels'] = upstream_shape['out_channels'] - except (MissingShapeDataError, ShapeInferenceError) as e: - logger.warning(f"Shape inference warning for node {node_id}: {e}. Using default.") - errors.append(e) - shape_info['in_channels'] = 3 - else: - shape_info['in_channels'] = 3 - - # Output channels: use metadata if available, otherwise config - if 'out_channels' not in shape_info: - shape_info['out_channels'] = config.get('filters', 64) - - # Spatial dimensions: use metadata if available, otherwise calculate - if 'out_height' not in shape_info or 'out_width' not in shape_info: - if incoming and incoming[0] in shape_map: - try: - prev_shape = safe_get_shape_data( - shape_map=shape_map, - node_id=node_id, - upstream_node_id=incoming[0], - required_keys=['out_height', 'out_width'], - default_values=None - ) - kernel_size = config.get('kernel_size', 3) - strides = config.get('strides', 1) - padding = config.get('padding', 'valid') - - if padding == 'same': - # Same padding preserves dimensions (with stride) - shape_info['out_height'] = (prev_shape['out_height'] + strides - 1) // strides - shape_info['out_width'] = (prev_shape['out_width'] + strides - 1) // strides - else: # valid padding - shape_info['out_height'] = (prev_shape['out_height'] - kernel_size) // strides + 1 - shape_info['out_width'] = (prev_shape['out_width'] - kernel_size) // strides + 1 - except (MissingShapeDataError, ShapeInferenceError) as e: - logger.warning(f"Could not compute spatial dimensions for conv2d {node_id}: {e}") - errors.append(e) - - elif node_type in ('maxpool2d', 'maxpool'): - # MaxPool preserves channels from upstream - if incoming and incoming[0] in shape_map: - try: - prev_shape = safe_get_shape_data( - shape_map=shape_map, - node_id=node_id, - upstream_node_id=incoming[0], - required_keys=['out_channels'], - default_values={'out_channels': 64} - ) - shape_info['out_channels'] = prev_shape['out_channels'] - except (MissingShapeDataError, ShapeInferenceError) as e: - logger.warning(f"Shape inference warning for maxpool {node_id}: {e}") - errors.append(e) - shape_info['out_channels'] = 64 - else: - shape_info['out_channels'] = 64 - - # Spatial dimensions: use metadata if available, otherwise calculate - if 'out_height' not in shape_info or 'out_width' not in shape_info: - if incoming and incoming[0] in shape_map: - try: - prev_shape = safe_get_shape_data( - shape_map=shape_map, - node_id=node_id, - upstream_node_id=incoming[0], - required_keys=['out_height', 'out_width'], - default_values={'out_height': 7, 'out_width': 7} - ) - pool_size = config.get('pool_size', 2) - strides = config.get('strides', 2) - padding = config.get('padding', 'valid') - - if padding == 'same': - shape_info['out_height'] = (prev_shape['out_height'] + strides - 1) // strides - shape_info['out_width'] = (prev_shape['out_width'] + strides - 1) // strides - else: # valid padding - shape_info['out_height'] = (prev_shape['out_height'] - pool_size) // strides + 1 - shape_info['out_width'] = (prev_shape['out_width'] - pool_size) // strides + 1 - except (MissingShapeDataError, ShapeInferenceError) as e: - logger.warning(f"Could not compute spatial dimensions for maxpool {node_id}: {e}") - errors.append(e) - - elif node_type == 'flatten': - # Flatten converts spatial dimensions to units - # Use metadata if available, otherwise calculate from upstream - if 'out_units' not in shape_info: - if incoming and incoming[0] in shape_map: - try: - prev_shape = safe_get_shape_data( - shape_map=shape_map, - node_id=node_id, - upstream_node_id=incoming[0], - required_keys=['out_channels', 'out_height', 'out_width'], - default_values={'out_channels': 64, 'out_height': 7, 'out_width': 7} - ) - channels = prev_shape['out_channels'] - height = prev_shape['out_height'] - width = prev_shape['out_width'] - shape_info['out_units'] = channels * height * width - except (MissingShapeDataError, ShapeInferenceError) as e: - logger.warning(f"Shape inference warning for flatten {node_id}: {e}") - errors.append(e) - shape_info['out_units'] = 3136 # 64 * 7 * 7 - else: - shape_info['out_units'] = 3136 # Default - - elif node_type == 'linear': - # Get input units from upstream layer (ALWAYS required) - if incoming and incoming[0] in shape_map: - try: - upstream_shape = safe_get_shape_data( - shape_map=shape_map, - node_id=node_id, - upstream_node_id=incoming[0], - required_keys=['out_units'], - default_values={'out_units': 512} - ) - shape_info['in_units'] = upstream_shape['out_units'] - except (MissingShapeDataError, ShapeInferenceError) as e: - logger.warning(f"Shape inference warning for linear {node_id}: {e}") - errors.append(e) - shape_info['in_units'] = 512 - else: - shape_info['in_units'] = 512 - - # Output units: use metadata if available, otherwise config - if 'out_units' not in shape_info: - shape_info['out_units'] = config.get('units', 128) - - elif node_type in ('batchnorm', 'batchnorm2d'): - # BatchNorm preserves all dimensions from upstream - # Only copy upstream if metadata doesn't provide them - if not metadata_shape and incoming and incoming[0] in shape_map: - try: - prev_shape = safe_get_shape_data( - shape_map=shape_map, - node_id=node_id, - upstream_node_id=incoming[0], - required_keys=[], # Accept whatever keys exist - default_values={} - ) - shape_info.update(prev_shape) - except (MissingShapeDataError, ShapeInferenceError) as e: - logger.warning(f"Shape inference warning for batchnorm {node_id}: {e}") - errors.append(e) - - elif node_type == 'group': - # Group blocks: Use metadata if available, otherwise compute from internal structure - if not metadata_shape: - # No metadata - compute output shape using shape computer - if shape_computer: - group_def_id = node.get('data', {}).get('groupDefinitionId') - - if group_def_id and incoming and incoming[0] in shape_map: - # Get input shape from upstream node - input_shape = shape_map[incoming[0]] - - # Compute output shape using internal structure - logger.debug(f"Computing shape for group block {node_id} (def: {group_def_id})") - output_shape, shape_errors = shape_computer.compute_output_shape( - group_def_id, - input_shape - ) - - # Collect any errors from shape computation - errors.extend(shape_errors) - - if output_shape: - shape_info = output_shape - logger.debug(f"Group block {node_id} output shape: {output_shape}") - else: - # Fallback: copy input shape - shape_info = input_shape.copy() - logger.warning(f"Failed to compute shape for group block {node_id}, using input shape") - elif incoming and incoming[0] in shape_map: - # No definition found, copy input shape - shape_info = shape_map[incoming[0]].copy() - logger.warning(f"Group block {node_id} has no definition ID, using input shape") - else: - # No input, use default - shape_info = {'out_channels': 3, 'out_height': 224, 'out_width': 224} - logger.warning(f"Group block {node_id} has no incoming edges, using default shape") - else: - # No shape computer available, fall back to old behavior - if incoming and incoming[0] in shape_map: - prev_shape = shape_map[incoming[0]] - # Copy input shape as default - shape_info.update(prev_shape) - else: - # Default starting shape - shape_info['out_channels'] = 3 - shape_info['out_height'] = 224 - shape_info['out_width'] = 224 - - else: - # For other layers: Use metadata if available, otherwise preserve upstream shape - if not metadata_shape and incoming and incoming[0] in shape_map: - prev_shape = shape_map[incoming[0]] - shape_info.update(prev_shape) - - shape_map[node_id] = shape_info - - return shape_map, errors - - -def validate_shape_map( - nodes: List[Dict], - shape_map: Dict[str, Dict[str, Any]] -) -> List[Exception]: - """ - Validate computed shape map for common critical issues (TensorFlow version). - - This catches problems that would cause runtime errors in generated code: - - Missing shape information - - Invalid dimensions (zero or negative) - - Type-specific requirements not met - - Args: - nodes: List of all nodes - shape_map: Computed shape mapping - - Returns: - List of validation errors (as exceptions for consistency with shape_errors) - """ - errors = [] - - for node in nodes: - node_id = node['id'] - node_type = get_node_type(node) - - # Skip non-layer nodes - if node_type in ('input', 'output', 'dataloader', 'group'): - continue - - shape_info = shape_map.get(node_id) - - # Critical: Shape info must exist - if not shape_info: - errors.append(ShapeInferenceError( - node_id=node_id, - node_type=node_type, - reason="No shape information computed for node", - suggestion="Check that node has valid upstream connections and metadata" - )) - continue - - # Type-specific validation - if node_type == 'linear' or node_type == 'dense': - # Linear/Dense MUST have in_features or in_units - if 'in_features' not in shape_info and 'in_units' not in shape_info: - errors.append(ShapeInferenceError( - node_id=node_id, - node_type=node_type, - reason="Missing required in_features/in_units for Linear/Dense layer", - suggestion="Check upstream Flatten or Linear layer output shape" - )) - # in_features/in_units must be positive - in_val = shape_info.get('in_features') or shape_info.get('in_units', 0) - if in_val <= 0: - errors.append(ShapeInferenceError( - node_id=node_id, - node_type=node_type, - reason=f"Invalid in_features/in_units={in_val} (must be > 0)", - suggestion="Check upstream layer produces valid output shape" - )) - - elif node_type == 'conv2d': - # Conv2d MUST have in_channels - if 'in_channels' not in shape_info: - errors.append(ShapeInferenceError( - node_id=node_id, - node_type=node_type, - reason="Missing required in_channels for Conv2d layer", - suggestion="Check upstream Conv2d or Input layer provides channels" - )) - - elif node_type == 'flatten': - # Flatten MUST produce out_features - if 'out_features' not in shape_info: - errors.append(ShapeInferenceError( - node_id=node_id, - node_type=node_type, - reason="Flatten layer must produce out_features", - suggestion="Check upstream layer has spatial dimensions (NHWC format)" - )) - elif shape_info.get('out_features', 0) <= 0: - errors.append(ShapeInferenceError( - node_id=node_id, - node_type=node_type, - reason=f"Invalid out_features={shape_info.get('out_features')} (must be > 0)", - suggestion="Check upstream layer output dimensions are valid" - )) - - return errors - - -def collect_all_nodes_with_internals( - main_nodes: List[Dict], - block_generator: Optional[TensorFlowBlockGenerator] = None -) -> List[Tuple[Dict, int, str]]: - """ - Collect all nodes including internal nodes from group blocks. - Returns list of tuples: (node, index, source_context) - source_context is either 'main' or 'group_{group_def_id}' - - This ensures we generate layer classes for ALL nodes, not just main model nodes. - """ - all_nodes = [] - node_index = 0 - - # Add main model nodes - for node in main_nodes: - all_nodes.append((node, node_index, 'main')) - node_index += 1 - - # Add internal nodes from group definitions - if block_generator: - for group_def_id, group_def in block_generator.group_definitions.items(): - internal_structure = group_def.get('internal_structure', {}) - internal_nodes = internal_structure.get('nodes', []) - - for internal_node in internal_nodes: - node_type = get_node_type(internal_node) - # Skip input/output nodes - if node_type not in ('input', 'dataloader', 'output'): - all_nodes.append((internal_node, node_index, f'group_{group_def_id}')) - node_index += 1 - - return all_nodes - - -def get_layer_signature(node: Dict, config: Dict[str, Any], node_type: str) -> str: - """ - Generate a unique signature for a layer based on its type and config. - Used for deduplication - layers with same signature can share the same class. - """ - if node_type == 'conv2d': - return f"conv2d_{config.get('out_channels', 64)}_{config.get('kernel_size', 3)}_{config.get('stride', 1)}_{config.get('padding', 0)}_{config.get('dilation', 1)}" - elif node_type == 'linear': - return f"linear_{config.get('out_features', 128)}_{config.get('bias', True)}" - elif node_type == 'maxpool': - return f"maxpool_{config.get('kernel_size', 2)}_{config.get('stride', 2)}_{config.get('padding', 0)}" - elif node_type == 'dropout': - return f"dropout_{config.get('p', 0.5)}" - elif node_type == 'batchnorm': - return f"batchnorm_{config.get('eps', 1e-5)}_{config.get('momentum', 0.1)}_{config.get('affine', True)}" - elif node_type == 'softmax': - return f"softmax_{config.get('dim', 1)}" - elif node_type == 'attention': - return f"attention_{config.get('embed_dim', 512)}_{config.get('num_heads', 8)}_{config.get('dropout', 0.0)}" - elif node_type == 'custom': - return f"custom_{config.get('name', 'CustomLayer')}" - else: - # For layers without config (relu, flatten, etc.) - return node_type - - -def generate_model_file( - nodes: List[Dict], - edges: List[Dict], - project_name: str, - shape_map: Dict[str, Dict[str, Any]], - block_generator: Optional[TensorFlowBlockGenerator] = None -) -> str: - """Generate complete model.py file with layer classes and main model class""" - - class_name = to_class_name(project_name) - - # Generate block class definitions FIRST (if any) - this populates the cache - block_classes_code = "" - if block_generator: - block_classes_code = block_generator.generate_all_block_classes() - - # COLLECT ALL NODES (main + internal from groups) and generate layer classes - all_nodes_to_generate = collect_all_nodes_with_internals(nodes, block_generator) - - # DEDUPLICATE by signature and generate layer classes - seen_signatures = set() - layer_classes = [] - - for node, idx, source_context in all_nodes_to_generate: - node_type = get_node_type(node) - config = node.get('data', {}).get('config', {}) - node_id = node['id'] - - # Get shape info (use shape_map for main nodes, extract for internal) - if source_context == 'main': - shape_info = shape_map.get(node_id, {}) - else: - shape_info = extract_shape_info_from_node(node) - - # Generate signature for deduplication - signature = get_layer_signature(node, config, node_type) - - # Only generate if we haven't seen this signature before - if signature not in seen_signatures: - seen_signatures.add(signature) - layer_class_code = generate_layer_class(node, idx, config, node_type, shape_info) - if layer_class_code: - layer_classes.append(layer_class_code) - - # Now generate layer instantiations and forward pass for MAIN MODEL ONLY - layer_instantiations = [] - forward_pass_lines = [] - - # Build edge map for forward pass - edge_map = {} - for edge in edges: - target = edge.get('target') - source = edge.get('source') - if target not in edge_map: - edge_map[target] = [] - edge_map[target].append(source) - - var_map = {} # Map node_id to variable name - - for idx, node in enumerate(nodes): - node_id = node['id'] - node_type = get_node_type(node) - config = node.get('data', {}).get('config', {}) - shape_info = shape_map.get(node_id, {}) - - if node_type in ('input', 'dataloader', 'output'): - # Skip input/output nodes - var_map[node_id] = 'x' if not var_map else 'x' - continue - - # Handle group blocks differently - if node_type == 'group': - # Get the group definition ID - group_def_id = node.get('data', {}).get('groupDefinitionId') - - if block_generator and group_def_id: - # Use the block class name from the generator - block_class_name = block_generator.get_block_class_name(group_def_id) - - if block_class_name: - layer_name = f"block_{node_id.replace('-', '_')}" - - # Get upstream node's output shape from shape_map - incoming = edge_map.get(node_id, []) - params = [] - - if incoming and incoming[0] in shape_map: - # Get upstream node's output shape - upstream_shape = shape_map[incoming[0]] - - # Extract in_channels or in_features from upstream shape - # TensorFlow uses same parameter names as PyTorch for consistency - # Pass in_channels if the upstream outputs channels (convolutional layers) - if 'out_channels' in upstream_shape: - in_channels = upstream_shape['out_channels'] - params.append(f"in_channels={in_channels}") - logger.debug(f"TF Block {node_id}: passing in_channels={in_channels} from upstream node {incoming[0]}") - - # Pass in_features if the upstream outputs features (linear layers) - # TensorFlow uses 'out_units' instead of 'out_features' - elif 'out_units' in upstream_shape: - in_units = upstream_shape['out_units'] - params.append(f"in_features={in_units}") - logger.debug(f"TF Block {node_id}: passing in_features={in_units} from upstream node {incoming[0]}") - elif 'out_features' in upstream_shape: - in_features = upstream_shape['out_features'] - params.append(f"in_features={in_features}") - logger.debug(f"TF Block {node_id}: passing in_features={in_features} from upstream node {incoming[0]}") - - # Pass num_features if the upstream outputs num_features (batch norm) - elif 'num_features' in upstream_shape: - num_features = upstream_shape['num_features'] - params.append(f"num_features={num_features}") - logger.debug(f"TF Block {node_id}: passing num_features={num_features} from upstream node {incoming[0]}") - else: - # Upstream shape exists but doesn't have expected keys - logger.warning(f"TF Block {node_id}: upstream shape {upstream_shape} doesn't contain expected keys") - else: - # Handle case where no upstream exists (use input node shape) - # Look for input nodes in the graph - input_nodes = [n for n in nodes if get_node_type(n) == 'input'] - if input_nodes and input_nodes[0]['id'] in shape_map: - input_shape = shape_map[input_nodes[0]['id']] - - # Use input node's output shape - if 'out_channels' in input_shape: - in_channels = input_shape['out_channels'] - params.append(f"in_channels={in_channels}") - logger.debug(f"TF Block {node_id}: no upstream, using input shape in_channels={in_channels}") - elif 'out_units' in input_shape: - in_units = input_shape['out_units'] - params.append(f"in_features={in_units}") - logger.debug(f"TF Block {node_id}: no upstream, using input shape in_features={in_units}") - elif 'out_features' in input_shape: - in_features = input_shape['out_features'] - params.append(f"in_features={in_features}") - logger.debug(f"TF Block {node_id}: no upstream, using input shape in_features={in_features}") - else: - logger.warning(f"TF Block {node_id}: input shape {input_shape} doesn't contain expected keys") - else: - # No upstream and no input node, use defaults - logger.warning(f"TF Block {node_id}: no upstream connection and no input node found") - - # Generate instantiation with computed parameters - # Each instance gets independent shape computation based on its position in the graph - if params: - layer_instantiations.append(f"self.{layer_name} = {block_class_name}({', '.join(params)}) # Instance at position {idx}") - else: - layer_instantiations.append(f"self.{layer_name} = {block_class_name}() # Instance at position {idx}") - - # Generate forward pass line - input_var = get_input_variable(incoming, var_map) - output_var = 'x' - forward_pass_lines.append(f"{output_var} = self.{layer_name}({input_var}, training=training)") - var_map[node_id] = output_var - else: - # Block class not found, skip - logger.warning(f"TF Block class not found for group definition {group_def_id}") - var_map[node_id] = 'x' - else: - # No block generator or definition ID, skip - logger.warning(f"TF No block generator or definition ID for node {node_id}") - var_map[node_id] = 'x' - continue - - # For regular nodes, we already generated the layer class above (no need to generate again) - - # Generate layer instantiation for __init__ - layer_name = get_layer_variable_name(node_type, idx, config) - layer_class_name = get_layer_class_name(node_type, idx, config) - layer_init = generate_layer_instantiation(layer_class_name, layer_name, shape_info) - if layer_init: - layer_instantiations.append(layer_init) - - # Generate forward pass line - incoming = edge_map.get(node_id, []) - input_var = get_input_variable(incoming, var_map) - output_var = 'x' - - forward_line = generate_forward_line(node_type, layer_name, input_var, output_var, shape_info) - if forward_line: - forward_pass_lines.append(forward_line) - - var_map[node_id] = output_var - - # Assemble the complete file - code = f'''""" -Generated TensorFlow/Keras Model -Architecture: {class_name} -Generated by VisionForge - -Note: TensorFlow uses NHWC format (batch, height, width, channels) -""" - -import tensorflow as tf -from tensorflow import keras -from tensorflow.keras import layers -from typing import Tuple, Optional - - -''' - - # Add block class definitions (already generated at the start) - if block_classes_code: - code += block_classes_code + '\n\n' - - # Add all layer class definitions - for layer_class in layer_classes: - code += layer_class + '\n\n' - - # Add main model class - code += f''' -class {class_name}(keras.Model): - """ - Main model class combining all layers. - - This model was automatically generated from a visual architecture. - Each layer is implemented as a separate class for clarity and reusability. - - Note: TensorFlow uses NHWC format (batch, height, width, channels) - """ - - def __init__(self): - """Initialize all layers in the model.""" - super({class_name}, self).__init__() - -''' - - # Add layer instantiations - for init in layer_instantiations: - code += f' {init}\n' - - code += ''' - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the model. - - Args: - inputs: Input tensor in NHWC format - training: Whether the model is in training mode - - Returns: - Output tensor after passing through all layers - """ - x = inputs -''' - - # Add forward pass lines - for line in forward_pass_lines: - code += f' {line}\n' - - code += ''' - return x - - -def create_model() -> keras.Model: - """ - Create and return an instance of the model. - - Returns: - Initialized model ready for training or inference - """ - model = {class_name}() - return model - - -if __name__ == '__main__': - # Create model - model = create_model() - print(f"Model: {class_name}") - - # Build the model with a sample input to initialize weights - model.build(input_shape=(None, 224, 224, 3)) # NHWC format - - # Print model summary - model.summary() - - # Test forward pass with dummy input - dummy_input = tf.random.normal([1, 224, 224, 3]) # NHWC format - output = model(dummy_input) - print(f"\\nInput shape: {{dummy_input.shape}}") # NHWC: [batch, height, width, channels] - print(f"Output shape: {{output.shape}}") -'''.format(class_name=class_name) - - return code - - -def generate_layer_class( - node: Dict, - idx: int, - config: Dict[str, Any], - node_type: str, - shape_info: Dict[str, Any] -) -> Optional[str]: - """Generate a complete layer class definition with documentation""" - - # Special node types that don't generate individual layer classes: - # - input/output/dataloader: Architectural markers for graph structure - # - group: Reusable components generated separately by BlockGenerator - if node_type in ('input', 'output', 'dataloader', 'group'): - return None - - class_name = get_layer_class_name(node_type, idx, config) - - if node_type == 'conv2d': - filters = config.get('filters', 64) - kernel_size = config.get('kernel_size', 3) - strides = config.get('strides', 1) - padding = config.get('padding', 'valid') - activation = config.get('activation', 'None') - activation_str = f"'{activation}'" if activation != 'None' else 'None' - - # Calculate output shape - out_h = shape_info.get('out_height', '?') - out_w = shape_info.get('out_width', '?') - out_c = filters - - return f'''class {class_name}(layers.Layer): - """ - 2D Convolutional Layer - - Applies a 2D convolution over an input signal. - - Parameters: - - Filters (output channels): {filters} - - Kernel size: {kernel_size}x{kernel_size} - - Strides: {strides} - - Padding: '{padding}' - - Activation: {activation if activation != 'None' else 'None'} - - Shape: - - Input: [batch, H, W, C] (NHWC format) - - Output: [batch, {out_h}, {out_w}, {out_c}] - """ - - def __init__(self): - """Initialize the convolutional layer.""" - super({class_name}, self).__init__() - self.conv = layers.Conv2D( - filters={filters}, - kernel_size={kernel_size}, - strides={strides}, - padding='{padding}', - activation={activation_str} - ) - - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the convolutional layer. - - Args: - inputs: Input tensor of shape [batch, H, W, C] - training: Whether in training mode - - Returns: - Output tensor of shape [batch, {out_h}, {out_w}, {out_c}] - """ - # Apply convolution - x = self.conv(inputs) - return x''' - - elif node_type == 'linear': - units = config.get('units', 128) - activation = config.get('activation', 'None') - use_bias = config.get('use_bias', True) - activation_str = f"'{activation}'" if activation != 'None' else 'None' - - return f'''class {class_name}(layers.Layer): - """ - Fully Connected (Dense) Layer - - Applies a linear transformation to the incoming data: y = xW + b - - Parameters: - - Units (output size): {units} - - Activation: {activation if activation != 'None' else 'None'} - - Use bias: {use_bias} - - Shape: - - Input: [batch, input_dim] - - Output: [batch, {units}] - """ - - def __init__(self): - """Initialize the dense layer.""" - super({class_name}, self).__init__() - self.dense = layers.Dense( - units={units}, - activation={activation_str}, - use_bias={use_bias} - ) - - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the dense layer. - - Args: - inputs: Input tensor of shape [batch, input_dim] - training: Whether in training mode - - Returns: - Output tensor of shape [batch, {units}] - """ - # Apply linear transformation - x = self.dense(inputs) - return x''' - - elif node_type in ('maxpool2d', 'maxpool'): - pool_size = config.get('pool_size', 2) - strides = config.get('strides', 2) - padding = config.get('padding', 'valid') - - return f'''class {class_name}(layers.Layer): - """ - 2D Max Pooling Layer - - Applies a 2D max pooling over an input signal. - Reduces spatial dimensions while preserving channel count. - - Parameters: - - Pool size: {pool_size}x{pool_size} - - Strides: {strides} - - Padding: '{padding}' - - Shape: - - Input: [batch, H, W, C] (NHWC format) - - Output: [batch, H/{strides}, W/{strides}, C] - """ - - def __init__(self): - """Initialize the max pooling layer.""" - super({class_name}, self).__init__() - self.pool = layers.MaxPooling2D( - pool_size={pool_size}, - strides={strides}, - padding='{padding}' - ) - - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the pooling layer. - - Args: - inputs: Input tensor of shape [batch, H, W, C] - training: Whether in training mode - - Returns: - Output tensor with reduced spatial dimensions - """ - # Apply max pooling - x = self.pool(inputs) - return x''' - - elif node_type == 'flatten': - out_units = shape_info.get('out_units', '?') - - return f'''class {class_name}(layers.Layer): - """ - Flatten Layer - - Flattens the input tensor to a 1D vector per batch sample. - Commonly used to transition from convolutional layers to fully connected layers. - - Shape: - - Input: [batch, H, W, C] (NHWC format) - - Output: [batch, H*W*C] = [batch, {out_units}] - """ - - def __init__(self): - """Initialize the flatten layer.""" - super({class_name}, self).__init__() - self.flatten = layers.Flatten() - - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the flatten layer. - - Args: - inputs: Input tensor of shape [batch, H, W, C] - training: Whether in training mode - - Returns: - Output tensor of shape [batch, H*W*C] - """ - # Flatten spatial and channel dimensions - x = self.flatten(inputs) - return x''' - - elif node_type == 'dropout': - rate = config.get('rate', 0.5) - - return f'''class {class_name}(layers.Layer): - """ - Dropout Regularization Layer - - Randomly sets input units to 0 with frequency rate during training. - Helps prevent overfitting. - - Parameters: - - Dropout rate: {rate} - - Shape: - - Input: [batch, *] (any shape) - - Output: [batch, *] (same shape as input) - """ - - def __init__(self): - """Initialize the dropout layer.""" - super({class_name}, self).__init__() - self.dropout = layers.Dropout(rate={rate}) - - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the dropout layer. - - Args: - inputs: Input tensor - training: Whether in training mode (dropout only active during training) - - Returns: - Output tensor with dropout applied during training - """ - # Apply dropout (only active during training) - x = self.dropout(inputs, training=training) - return x''' - - elif node_type in ('batchnorm', 'batchnorm2d'): - momentum = config.get('momentum', 0.99) - epsilon = config.get('epsilon', 0.001) - - return f'''class {class_name}(layers.Layer): - """ - Batch Normalization Layer - - Normalizes the activations of the previous layer at each batch. - Helps stabilize and accelerate training. - - Parameters: - - Momentum: {momentum} - - Epsilon: {epsilon} - - Shape: - - Input: [batch, H, W, C] or [batch, features] (NHWC format) - - Output: Same shape as input - """ - - def __init__(self): - """Initialize the batch normalization layer.""" - super({class_name}, self).__init__() - self.bn = layers.BatchNormalization( - momentum={momentum}, - epsilon={epsilon} - ) - - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the batch normalization layer. - - Args: - inputs: Input tensor - training: Whether in training mode - - Returns: - Normalized output tensor of same shape - """ - # Apply batch normalization - x = self.bn(inputs, training=training) - return x''' - - elif node_type == 'concat': - axis = config.get('axis', -1) - - return f'''class {class_name}(layers.Layer): - """ - Concatenation Layer - - Concatenates multiple input tensors along a specified axis. - Used for skip connections and multi-path architectures. - - Parameters: - - Axis: {axis} - - Shape: - - Input: List of tensors with compatible shapes - - Output: Single concatenated tensor - """ - - def __init__(self): - """Initialize the concatenation layer.""" - super({class_name}, self).__init__() - self.concat = layers.Concatenate(axis={axis}) - - def call(self, inputs: list, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the concatenation layer. - - Args: - inputs: List of input tensors to concatenate - training: Whether in training mode - - Returns: - Concatenated output tensor - """ - # Concatenate along specified axis - x = self.concat(inputs) - return x''' - - elif node_type == 'add': - return f'''class {class_name}(layers.Layer): - """ - Addition Layer - - Performs element-wise addition of multiple input tensors. - Used for residual connections and multi-path architectures. - - Shape: - - Input: List of tensors with identical shapes - - Output: Single tensor with same shape as inputs - """ - - def __init__(self): - """Initialize the addition layer.""" - super({class_name}, self).__init__() - self.add = layers.Add() - - def call(self, inputs: list, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the addition layer. - - Args: - inputs: List of input tensors to add - training: Whether in training mode - - Returns: - Sum of input tensors - """ - # Element-wise addition - x = self.add(inputs) - return x''' - - elif node_type == 'custom': - name = config.get('name', 'CustomLayer') - description = config.get('description', 'User-defined custom layer') - - # Generate proper class name from user's layer name - safe_name = name.replace(' ', '_').replace('-', '_') - custom_class_name = f"CustomLayer_{safe_name}" - - return f'''class {custom_class_name}(layers.Layer): - """ - Custom User-Defined Layer: {name} - - {description} - - TODO: Implement your custom layer logic below. - This class provides the basic structure following TensorFlow/Keras conventions. - Add your initialization and call method logic. - - Note: TensorFlow uses NHWC format (batch, height, width, channels) - - Shape: - - Input: [batch, *] (Define your input shape in NHWC format) - - Output: [batch, *] (Define your output shape) - """ - - def __init__(self): - """Initialize the custom layer.""" - super({custom_class_name}, self).__init__() - - # TODO: Define your layer components here - # Examples: - # self.dense = layers.Dense(units=128) - # self.conv = layers.Conv2D(filters=64, kernel_size=3) - # self.activation = layers.ReLU() - # self.dropout = layers.Dropout(rate=0.5) - - pass - - def call(self, inputs: tf.Tensor, training: Optional[bool] = None) -> tf.Tensor: - """ - Forward pass through the custom layer. - - Args: - inputs: Input tensor in NHWC format - training: Whether in training mode - - Returns: - Output tensor - """ - # TODO: Implement your call method logic here - # Examples: - # x = self.dense(inputs) - # x = self.activation(x) - # x = self.dropout(x, training=training) - - # Placeholder: returns input unchanged - # Replace this with your custom logic - x = inputs - return x''' - - # If we reach here, the node type is not supported - raise UnsupportedNodeTypeError( - node_id=node.get('id', 'unknown'), - node_type=node_type, - framework='TensorFlow' - ) - - -def generate_layer_instantiation( - class_name: str, - layer_name: str, - shape_info: Dict[str, Any] -) -> str: - """Generate layer instantiation line for __init__ method""" - # TensorFlow layers don't need input size in constructor - if 'in_channels' in shape_info: - in_ch = shape_info['in_channels'] - return f"self.{layer_name} = {class_name}() # Input: {in_ch} channels (NHWC)" - elif 'in_units' in shape_info: - in_units = shape_info['in_units'] - return f"self.{layer_name} = {class_name}() # Input: {in_units} units" - else: - return f"self.{layer_name} = {class_name}()" - - -def generate_forward_line( - node_type: str, - layer_name: str, - input_var: str, - output_var: str, - shape_info: Dict[str, Any] -) -> str: - """Generate forward pass line with shape comments""" - # Build shape comment - shape_comment = "" - if 'out_channels' in shape_info: - h = shape_info.get('out_height', '?') - w = shape_info.get('out_width', '?') - c = shape_info['out_channels'] - shape_comment = f" # Shape: [batch, {h}, {w}, {c}] (NHWC)" - elif 'out_units' in shape_info: - u = shape_info['out_units'] - shape_comment = f" # Shape: [batch, {u}]" - - # Handle layers that need training parameter - if node_type in ('dropout', 'batchnorm', 'batchnorm2d'): - return f"{output_var} = self.{layer_name}({input_var}, training=training){shape_comment}" - # Handle merge layers - elif node_type in ('concat', 'add'): - return f"{output_var} = self.{layer_name}({input_var}){shape_comment}" - else: - return f"{output_var} = self.{layer_name}({input_var}){shape_comment}" - - -def get_layer_class_name(node_type: str, idx: int, config: Dict[str, Any]) -> str: - """Generate descriptive class name for layer""" - type_name = node_type.replace('_', '').replace('2d', '2D').replace('3d', '3D').title() - - # Add descriptive suffix based on config - if node_type == 'conv2d': - filters = config.get('filters', 64) - kernel = config.get('kernel_size', 3) - return f"{type_name}Layer_{filters}filters_{kernel}x{kernel}" - elif node_type == 'linear': - units = config.get('units', 128) - return f"DenseLayer_{units}units" - elif node_type in ('maxpool2d', 'maxpool'): - pool_size = config.get('pool_size', 2) - return f"MaxPool2DLayer_{pool_size}x{pool_size}" - else: - return f"{type_name}Layer_{idx}" - - -def get_layer_variable_name(node_type: str, idx: int, config: Dict[str, Any]) -> str: - """Generate descriptive variable name for layer instance""" - # Create readable names based on layer type - if node_type == 'conv2d': - filters = config.get('filters', 64) - return f"conv_{filters}filters" - elif node_type == 'linear': - units = config.get('units', 128) - return f"dense_{units}" - elif node_type in ('maxpool2d', 'maxpool'): - return f"maxpool_{idx}" - elif node_type == 'flatten': - return f"flatten" - elif node_type == 'dropout': - return f"dropout_{idx}" - elif node_type in ('batchnorm', 'batchnorm2d'): - return f"batchnorm_{idx}" - elif node_type == 'concat': - return f"concat_{idx}" - elif node_type == 'add': - return f"add_{idx}" - else: - return f"layer_{idx}" - - -def get_input_variable(incoming: List[str], var_map: Dict[str, str]) -> str: - """Determine input variable name based on incoming connections""" - if not incoming: - return 'x' - elif len(incoming) == 1: - return var_map.get(incoming[0], 'x') - else: - # Multiple inputs (for concat, add, etc.) - input_vars = [var_map.get(src, 'x') for src in incoming] - return f"[{', '.join(input_vars)}]" - - -def generate_training_script(project_name: str) -> str: - """Generate comprehensive training script with best practices""" - class_name = to_class_name(project_name) - - return f'''""" -Training Script for {class_name} -Generated by VisionForge -""" - -import tensorflow as tf -from tensorflow import keras -import numpy as np -from pathlib import Path -from typing import Tuple, Dict, Optional - -from model import create_model -from dataset import CustomDataset - - -def train_model( - num_epochs: int = 10, - batch_size: int = 32, - learning_rate: float = 0.001, - weight_decay: float = 1e-4, - use_gpu: bool = True -) -> keras.callbacks.History: - """ - Main training function. - - Args: - num_epochs: Number of training epochs - batch_size: Batch size for training - learning_rate: Initial learning rate - weight_decay: L2 regularization factor - use_gpu: Whether to use GPU if available - - Returns: - Training history object - """ - # Configure GPU - if use_gpu: - gpus = tf.config.list_physical_devices('GPU') - if gpus: - print(f'Found {{len(gpus)}} GPU(s)') - for gpu in gpus: - tf.config.experimental.set_memory_growth(gpu, True) - else: - print('No GPU found, using CPU') - else: - tf.config.set_visible_devices([], 'GPU') - print('Using CPU') - - # Set random seeds for reproducibility - tf.random.set_seed(42) - np.random.seed(42) - - # Create model - model = create_model() - print(f'\\nModel created: {{model.__class__.__name__}}') - - # Build model with sample input (NHWC format) - model.build(input_shape=(None, 224, 224, 3)) - model.summary() - - # TODO: Replace with your actual dataset - # Example: - # train_dataset = CustomDataset('path/to/train', batch_size=batch_size) - # val_dataset = CustomDataset('path/to/val', batch_size=batch_size) - - print('\\nCreating dummy datasets (replace with actual data)...') - # Dummy data (NHWC format: batch, height, width, channels) - train_data = np.random.randn(1000, 224, 224, 3).astype(np.float32) - train_labels = np.random.randint(0, 10, (1000,)) - val_data = np.random.randn(200, 224, 224, 3).astype(np.float32) - val_labels = np.random.randint(0, 10, (200,)) - - # Create TensorFlow datasets - train_dataset = tf.data.Dataset.from_tensor_slices((train_data, train_labels)) - train_dataset = train_dataset.shuffle(1000).batch(batch_size).prefetch(tf.data.AUTOTUNE) - - val_dataset = tf.data.Dataset.from_tensor_slices((val_data, val_labels)) - val_dataset = val_dataset.batch(batch_size).prefetch(tf.data.AUTOTUNE) - - # Compile model - model.compile( - optimizer=keras.optimizers.Adam( - learning_rate=learning_rate, - weight_decay=weight_decay - ), - loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True), - metrics=['accuracy'] - ) - - # Callbacks - callbacks = [ - keras.callbacks.ModelCheckpoint( - 'best_model.keras', - save_best_only=True, - monitor='val_loss', - mode='min', - verbose=1 - ), - keras.callbacks.EarlyStopping( - monitor='val_loss', - patience=5, - restore_best_weights=True, - verbose=1 - ), - keras.callbacks.ReduceLROnPlateau( - monitor='val_loss', - factor=0.5, - patience=3, - min_lr=1e-7, - verbose=1 - ), - keras.callbacks.TensorBoard( - log_dir='./logs', - histogram_freq=1 - ) - ] - - print(f'\\nStarting training for {{num_epochs}} epochs...\\n') - - # Train the model - history = model.fit( - train_dataset, - validation_data=val_dataset, - epochs=num_epochs, - callbacks=callbacks, - verbose=1 - ) - - # Save final model - model.save('{project_name}_final.keras') - print(f"\\nFinal model saved to {project_name}_final.keras") - - # Print training summary - print('\\n' + '=' * 60) - print('Training completed!') - print(f'Best validation loss: {{min(history.history["val_loss"]):.4f}}') - print(f'Best validation accuracy: {{max(history.history["val_accuracy"]):.4f}}') - print('=' * 60) - - return history - - -if __name__ == '__main__': - # Train the model - history = train_model( - num_epochs=10, - batch_size=32, - learning_rate=0.001, - weight_decay=1e-4, - use_gpu=True - ) - - print('\\nTraining complete!') -''' - - -def generate_dataset_class(nodes: List[Dict]) -> str: - """Generate dataset class for data loading""" - - return '''""" -Custom Dataset Class for TensorFlow -Generated by VisionForge -""" - -import tensorflow as tf -from tensorflow import keras -import numpy as np -from pathlib import Path -from typing import Tuple, Optional - - -class CustomDataset(keras.utils.PyDataset): - """ - Custom dataset using tf.keras.utils.PyDataset for efficient data loading. - - This is a template - replace with your actual data loading logic. - - Args: - data_path: Path to the dataset directory - batch_size: Number of samples per batch - shuffle: Whether to shuffle the data - split: Dataset split ('train', 'val', or 'test') - """ - - def __init__( - self, - data_path: str, - batch_size: int = 32, - shuffle: bool = True, - split: str = 'train', - **kwargs - ): - """ - Initialize the dataset. - - Args: - data_path: Path to data directory - batch_size: Batch size for loading - shuffle: Whether to shuffle data - split: Which split to load ('train', 'val', 'test') - """ - super().__init__(**kwargs) - self.data_path = Path(data_path) - self.batch_size = batch_size - self.shuffle = shuffle - self.split = split - - # TODO: Replace with your actual data loading - # Example: Load file paths and labels - # self.samples = self._load_samples() - - # For demonstration, create dummy data - self.num_samples = 1000 if split == 'train' else 200 - print(f'Loaded {{self.num_samples}} samples for {{split}} split') - - def __len__(self) -> int: - """Return number of batches per epoch.""" - return self.num_samples // self.batch_size - - def __getitem__(self, idx: int) -> Tuple[np.ndarray, np.ndarray]: - """ - Generate one batch of data. - - Args: - idx: Batch index - - Returns: - Tuple of (inputs, targets) in NHWC format - """ - # TODO: Replace with actual data loading - # Example: - # batch_samples = self.samples[idx*self.batch_size:(idx+1)*self.batch_size] - # batch_x = [] - # batch_y = [] - # for sample in batch_samples: - # image = load_image(sample['path']) # Load and preprocess - # batch_x.append(image) - # batch_y.append(sample['label']) - # return np.array(batch_x), np.array(batch_y) - - # Generate dummy batch (NHWC format: batch, height, width, channels) - batch_x = np.random.randn(self.batch_size, 224, 224, 3).astype(np.float32) - batch_y = np.random.randint(0, 10, self.batch_size) - - return batch_x, batch_y - - def on_epoch_end(self): - """Called at the end of each epoch.""" - if self.shuffle: - # TODO: Implement shuffling logic - pass - - def _load_samples(self): - """ - Load sample paths and labels from disk. - - Returns: - List of sample dictionaries with 'path' and 'label' keys - """ - # TODO: Implement actual data loading logic - # Example for image classification: - # - # samples = [] - # split_dir = self.data_path / self.split - # for class_idx, class_name in enumerate(sorted(split_dir.iterdir())): - # if class_name.is_dir(): - # for img_path in class_name.glob('*.jpg'): - # samples.append({{ - # 'path': str(img_path), - # 'label': class_idx - # }}) - # return samples - - pass - - -# Example data preprocessing functions -def preprocess_image(image_path: str, target_size: Tuple[int, int] = (224, 224)) -> np.ndarray: - """ - Load and preprocess an image. - - Args: - image_path: Path to the image file - target_size: Target size for resizing (height, width) - - Returns: - Preprocessed image array in NHWC format - """ - # Load image - image = tf.io.read_file(image_path) - image = tf.image.decode_jpeg(image, channels=3) - - # Resize - image = tf.image.resize(image, target_size) - - # Normalize to [0, 1] - image = tf.cast(image, tf.float32) / 255.0 - - # Optional: Normalize with ImageNet mean and std - # mean = tf.constant([0.485, 0.456, 0.406]) - # std = tf.constant([0.229, 0.224, 0.225]) - # image = (image - mean) / std - - return image.numpy() - - -def augment_image(image: np.ndarray) -> np.ndarray: - """ - Apply data augmentation to an image. - - Args: - image: Input image in NHWC format - - Returns: - Augmented image - """ - image = tf.constant(image) - - # Random horizontal flip - image = tf.image.random_flip_left_right(image) - - # Random brightness and contrast - image = tf.image.random_brightness(image, max_delta=0.2) - image = tf.image.random_contrast(image, lower=0.8, upper=1.2) - - # Random rotation (small angles) - # Note: Requires tf-addons for rotation - # image = tfa.image.rotate(image, angles=tf.random.uniform([], -0.2, 0.2)) - - # Clip values to [0, 1] - image = tf.clip_by_value(image, 0.0, 1.0) - - return image.numpy() - - -# Example usage -if __name__ == '__main__': - # Create dataset instances - train_dataset = CustomDataset('data/', batch_size=32, split='train') - val_dataset = CustomDataset('data/', batch_size=32, split='val') - - print(f'Train dataset: {{len(train_dataset)}} batches') - print(f'Val dataset: {{len(val_dataset)}} batches') - - # Get a sample batch - batch_x, batch_y = train_dataset[0] - print(f'\\nBatch X shape: {{batch_x.shape}}') # Should be (32, 224, 224, 3) in NHWC format - print(f'Batch Y shape: {{batch_y.shape}}') -''' - - -def generate_config_file(nodes: List[Dict]) -> str: - """Generate configuration file with hyperparameters""" - - # Find input shape from nodes (NHWC format) - input_shape = "[1, 224, 224, 3]" - for node in nodes: - if get_node_type(node) in ('input', 'dataloader'): - shape = node.get('data', {}).get('outputShape', {}).get('dims') - if shape: - input_shape = str(shape) - break - - return f'''""" -Configuration File -Generated by VisionForge -Contains all hyperparameters and settings for training - -Note: TensorFlow uses NHWC format (batch, height, width, channels) -""" - -# Training Configuration -BATCH_SIZE = 32 -LEARNING_RATE = 0.001 -NUM_EPOCHS = 10 -WEIGHT_DECAY = 1e-4 - -# Model Configuration (NHWC format: batch, height, width, channels) -INPUT_SHAPE = {input_shape} -NUM_CLASSES = 10 # TODO: Set to your number of classes - -# Optimizer Settings -OPTIMIZER = 'adam' # Options: 'adam', 'sgd', 'rmsprop', 'adamw' -MOMENTUM = 0.9 # For SGD -BETAS = (0.9, 0.999) # For Adam/AdamW - -# Learning Rate Scheduler -USE_SCHEDULER = True -SCHEDULER_TYPE = 'reduce_on_plateau' # Options: 'reduce_on_plateau', 'exponential', 'cosine' -LR_PATIENCE = 3 # For ReduceLROnPlateau -LR_FACTOR = 0.5 # For ReduceLROnPlateau -DECAY_STEPS = 1000 # For ExponentialDecay -DECAY_RATE = 0.96 # For ExponentialDecay - -# Early Stopping -USE_EARLY_STOPPING = True -EARLY_STOPPING_PATIENCE = 5 - -# Data Augmentation (for training) -USE_AUGMENTATION = True -RANDOM_FLIP = True -RANDOM_ROTATION = True -RANDOM_ZOOM = True -RANDOM_BRIGHTNESS = True -RANDOM_CONTRAST = True - -# Augmentation parameters -ROTATION_RANGE = 15 -WIDTH_SHIFT_RANGE = 0.1 -HEIGHT_SHIFT_RANGE = 0.1 -ZOOM_RANGE = 0.1 - -# Normalization (ImageNet statistics) -NORMALIZE = True -MEAN = [0.485, 0.456, 0.406] -STD = [0.229, 0.224, 0.225] - -# Device Configuration -USE_GPU = True # Use GPU if available -MEMORY_GROWTH = True # Allow GPU memory to grow as needed - -# Mixed Precision Training (for faster training on modern GPUs) -USE_MIXED_PRECISION = False - -# Checkpointing -SAVE_BEST_ONLY = True -CHECKPOINT_DIR = './checkpoints' -SAVE_FREQUENCY = 1 # Save every N epochs - -# Logging -USE_TENSORBOARD = True -TENSORBOARD_DIR = './logs' -LOG_HISTOGRAMS = True - -# Data Loading -NUM_PARALLEL_CALLS = tf.data.AUTOTUNE if 'tf' in dir() else 4 -PREFETCH_BUFFER = tf.data.AUTOTUNE if 'tf' in dir() else 2 - -# Paths -DATA_DIR = './data' -TRAIN_DIR = DATA_DIR + '/train' -VAL_DIR = DATA_DIR + '/val' -TEST_DIR = DATA_DIR + '/test' - -# Model specific -DROPOUT_RATE = 0.5 -BATCH_NORM_MOMENTUM = 0.99 -BATCH_NORM_EPSILON = 0.001 - -# Import TensorFlow for AUTOTUNE constant -try: - import tensorflow as tf - NUM_PARALLEL_CALLS = tf.data.AUTOTUNE - PREFETCH_BUFFER = tf.data.AUTOTUNE -except ImportError: - pass -''' - - -def get_node_type(node: Dict) -> str: - """Extract node type from node dictionary""" - return node.get('data', {}).get('blockType', node.get('type', 'unknown')) - + # Delegate to new template-based orchestrator + orchestrator = TensorFlowCodeOrchestrator() + return orchestrator.generate(nodes, edges, project_name, group_definitions) -def to_class_name(name: str) -> str: - """Convert project name to valid Python class name""" - import re - # Remove special characters and convert to PascalCase - name = re.sub(r'[^a-zA-Z0-9]', ' ', name) - name = ''.join(word.capitalize() for word in name.split()) - if not name: - return 'GeneratedModel' - if name[0].isdigit(): - name = 'Model' + name - return name +# ==================== LEGACY CODE REMOVED ==================== +# All legacy code after the return statement has been removed. +# The new implementation is in: +# - codegen/tensorflow_orchestrator.py +# - codegen/tensorflow_group_generator.py +# ============================================================== diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/adaptiveavgpool2d.ts b/project/frontend/src/lib/nodes/definitions/pytorch/adaptiveavgpool2d.ts new file mode 100644 index 0000000..12edd2a --- /dev/null +++ b/project/frontend/src/lib/nodes/definitions/pytorch/adaptiveavgpool2d.ts @@ -0,0 +1,81 @@ +/** + * PyTorch AdaptiveAvgPool2D Layer Node Definition + */ + +import { NodeDefinition } from '../../base' +import { NodeMetadata, BackendFramework } from '../../contracts' +import { TensorShape, BlockConfig, ConfigField, BlockType } from '../../../types' + +export class AdaptiveAvgPool2DNode extends NodeDefinition { + readonly metadata: NodeMetadata = { + type: 'adaptiveavgpool2d', + label: 'AdaptiveAvgPool2D', + category: 'basic', + color: 'var(--color-primary)', + icon: 'Resize', + description: 'Adaptive average pooling to fixed output size', + framework: BackendFramework.PyTorch + } + + readonly configSchema: ConfigField[] = [ + { + name: 'output_size', + label: 'Output Size', + type: 'text', + default: '1', + description: 'Target output size (single number or [H, W])' + } + ] + + computeOutputShape(inputShape: TensorShape | undefined, config: BlockConfig): TensorShape | undefined { + if (!inputShape) { + return undefined + } + + if (inputShape.dims.length !== 4) { + return undefined + } + + const [batch, channels] = inputShape.dims as number[] + const outputSizeStr = String(config.output_size ?? '1') + + let outHeight: number, outWidth: number + + try { + if (outputSizeStr.includes(',') || outputSizeStr.includes('[')) { + const cleaned = outputSizeStr.replace(/[\[\]\(\)\s]/g, '') + const parts = cleaned.split(',') + outHeight = parseInt(parts[0]) + outWidth = parts.length > 1 ? parseInt(parts[1]) : outHeight + } else { + outHeight = outWidth = parseInt(outputSizeStr) + } + } catch { + outHeight = outWidth = 1 + } + + return { + dims: [batch, channels, outHeight, outWidth], + description: `AdaptiveAvgPool2D(${outHeight}x${outWidth})` + } + } + + validateIncomingConnection( + sourceNodeType: BlockType, + sourceOutputShape: TensorShape | undefined, + targetConfig: BlockConfig + ): string | undefined { + if (sourceNodeType === 'input' || sourceNodeType === 'dataloader') { + return undefined + } + + if (sourceNodeType === 'empty' || sourceNodeType === 'custom') { + return undefined + } + + return this.validateDimensions(sourceOutputShape, { + dims: 4, + description: '[batch, channels, height, width]' + }) + } +} diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/avgpool2d.ts b/project/frontend/src/lib/nodes/definitions/pytorch/avgpool2d.ts new file mode 100644 index 0000000..c6acda8 --- /dev/null +++ b/project/frontend/src/lib/nodes/definitions/pytorch/avgpool2d.ts @@ -0,0 +1,88 @@ +/** + * PyTorch AvgPool2D Layer Node Definition + */ + +import { NodeDefinition } from '../../base' +import { NodeMetadata, BackendFramework } from '../../contracts' +import { TensorShape, BlockConfig, ConfigField, BlockType } from '../../../types' + +export class AvgPool2DNode extends NodeDefinition { + readonly metadata: NodeMetadata = { + type: 'avgpool2d', + label: 'AvgPool2D', + category: 'basic', + color: 'var(--color-primary)', + icon: 'ArrowsInSimple', + description: 'Average pooling for 2D inputs', + framework: BackendFramework.PyTorch + } + + readonly configSchema: ConfigField[] = [ + { + name: 'kernel_size', + label: 'Kernel Size', + type: 'number', + default: 2, + min: 1, + description: 'Size of the pooling window' + }, + { + name: 'stride', + label: 'Stride', + type: 'number', + default: 2, + min: 1, + description: 'Stride of the pooling window' + }, + { + name: 'padding', + label: 'Padding', + type: 'number', + default: 0, + min: 0, + description: 'Zero padding on both sides' + } + ] + + computeOutputShape(inputShape: TensorShape | undefined, config: BlockConfig): TensorShape | undefined { + if (!inputShape) { + return undefined + } + + if (inputShape.dims.length !== 4) { + return undefined + } + + const [batch, channels, height, width] = inputShape.dims as number[] + const kernel = (config.kernel_size ?? 2) as number + const stride = (config.stride ?? 2) as number + const padding = (config.padding ?? 0) as number + + const outHeight = Math.floor((height + 2 * padding - kernel) / stride) + 1 + const outWidth = Math.floor((width + 2 * padding - kernel) / stride) + 1 + + return { + dims: [batch, channels, outHeight, outWidth], + description: `AvgPool2D(${kernel}x${kernel})` + } + } + + validateIncomingConnection( + sourceNodeType: BlockType, + sourceOutputShape: TensorShape | undefined, + targetConfig: BlockConfig + ): string | undefined { + if (sourceNodeType === 'input' || sourceNodeType === 'dataloader') { + return undefined + } + + if (sourceNodeType === 'empty' || sourceNodeType === 'custom') { + return undefined + } + + return this.validateDimensions(sourceOutputShape, { + dims: 4, + description: '[batch, channels, height, width]' + }) + } +} diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/conv1d.ts b/project/frontend/src/lib/nodes/definitions/pytorch/conv1d.ts new file mode 100644 index 0000000..59e4160 --- /dev/null +++ b/project/frontend/src/lib/nodes/definitions/pytorch/conv1d.ts @@ -0,0 +1,112 @@ +/** + * PyTorch Conv1D Layer Node Definition + */ + +import { NodeDefinition } from '../../base' +import { NodeMetadata, BackendFramework } from '../../contracts' +import { TensorShape, BlockConfig, ConfigField, BlockType } from '../../../types' + +export class Conv1DNode extends NodeDefinition { + readonly metadata: NodeMetadata = { + type: 'conv1d', + label: 'Conv1D', + category: 'advanced', + color: 'var(--color-purple)', + icon: 'WaveSquare', + description: '1D convolutional layer', + framework: BackendFramework.PyTorch + } + + readonly configSchema: ConfigField[] = [ + { + name: 'out_channels', + label: 'Output Channels', + type: 'number', + required: true, + min: 1, + description: 'Number of output channels' + }, + { + name: 'kernel_size', + label: 'Kernel Size', + type: 'number', + default: 3, + min: 1, + description: 'Size of the convolving kernel' + }, + { + name: 'stride', + label: 'Stride', + type: 'number', + default: 1, + min: 1, + description: 'Stride of the convolution' + }, + { + name: 'padding', + label: 'Padding', + type: 'number', + default: 0, + min: 0, + description: 'Zero padding on both sides' + }, + { + name: 'dilation', + label: 'Dilation', + type: 'number', + default: 1, + min: 1, + description: 'Spacing between kernel elements' + }, + { + name: 'bias', + label: 'Use Bias', + type: 'boolean', + default: true, + description: 'Add learnable bias' + } + ] + + computeOutputShape(inputShape: TensorShape | undefined, config: BlockConfig): TensorShape | undefined { + if (!inputShape || !config.out_channels) { + return undefined + } + + if (inputShape.dims.length !== 3) { + return undefined + } + + const [batch, _, length] = inputShape.dims as number[] + + const kernel = (config.kernel_size ?? 3) as number + const stride = (config.stride ?? 1) as number + const padding = (config.padding ?? 0) as number + const dilation = (config.dilation ?? 1) as number + + const outLength = Math.floor((length + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1 + + return { + dims: [batch, config.out_channels as number, outLength], + description: `Conv1D(${config.out_channels})` + } + } + + validateIncomingConnection( + sourceNodeType: BlockType, + sourceOutputShape: TensorShape | undefined, + targetConfig: BlockConfig + ): string | undefined { + if (sourceNodeType === 'input' || sourceNodeType === 'dataloader') { + return undefined + } + + if (sourceNodeType === 'empty' || sourceNodeType === 'custom') { + return undefined + } + + return this.validateDimensions(sourceOutputShape, { + dims: 3, + description: '[batch, channels, length]' + }) + } +} diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/conv3d.ts b/project/frontend/src/lib/nodes/definitions/pytorch/conv3d.ts new file mode 100644 index 0000000..760f18e --- /dev/null +++ b/project/frontend/src/lib/nodes/definitions/pytorch/conv3d.ts @@ -0,0 +1,114 @@ +/** + * PyTorch Conv3D Layer Node Definition + */ + +import { NodeDefinition } from '../../base' +import { NodeMetadata, BackendFramework } from '../../contracts' +import { TensorShape, BlockConfig, ConfigField, BlockType } from '../../../types' + +export class Conv3DNode extends NodeDefinition { + readonly metadata: NodeMetadata = { + type: 'conv3d', + label: 'Conv3D', + category: 'advanced', + color: 'var(--color-purple)', + icon: 'Cube', + description: '3D convolutional layer', + framework: BackendFramework.PyTorch + } + + readonly configSchema: ConfigField[] = [ + { + name: 'out_channels', + label: 'Output Channels', + type: 'number', + required: true, + min: 1, + description: 'Number of output channels' + }, + { + name: 'kernel_size', + label: 'Kernel Size', + type: 'number', + default: 3, + min: 1, + description: 'Size of the convolving kernel' + }, + { + name: 'stride', + label: 'Stride', + type: 'number', + default: 1, + min: 1, + description: 'Stride of the convolution' + }, + { + name: 'padding', + label: 'Padding', + type: 'number', + default: 0, + min: 0, + description: 'Zero padding on all sides' + }, + { + name: 'dilation', + label: 'Dilation', + type: 'number', + default: 1, + min: 1, + description: 'Spacing between kernel elements' + }, + { + name: 'bias', + label: 'Use Bias', + type: 'boolean', + default: true, + description: 'Add learnable bias' + } + ] + + computeOutputShape(inputShape: TensorShape | undefined, config: BlockConfig): TensorShape | undefined { + if (!inputShape || !config.out_channels) { + return undefined + } + + if (inputShape.dims.length !== 5) { + return undefined + } + + const [batch, _, depth, height, width] = inputShape.dims as number[] + + const kernel = (config.kernel_size ?? 3) as number + const stride = (config.stride ?? 1) as number + const padding = (config.padding ?? 0) as number + const dilation = (config.dilation ?? 1) as number + + const outDepth = Math.floor((depth + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1 + const outHeight = Math.floor((height + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1 + const outWidth = Math.floor((width + 2 * padding - dilation * (kernel - 1) - 1) / stride) + 1 + + return { + dims: [batch, config.out_channels as number, outDepth, outHeight, outWidth], + description: `Conv3D(${config.out_channels})` + } + } + + validateIncomingConnection( + sourceNodeType: BlockType, + sourceOutputShape: TensorShape | undefined, + targetConfig: BlockConfig + ): string | undefined { + if (sourceNodeType === 'input' || sourceNodeType === 'dataloader') { + return undefined + } + + if (sourceNodeType === 'empty' || sourceNodeType === 'custom') { + return undefined + } + + return this.validateDimensions(sourceOutputShape, { + dims: 5, + description: '[batch, channels, depth, height, width]' + }) + } +} diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/embedding.ts b/project/frontend/src/lib/nodes/definitions/pytorch/embedding.ts new file mode 100644 index 0000000..967d20d --- /dev/null +++ b/project/frontend/src/lib/nodes/definitions/pytorch/embedding.ts @@ -0,0 +1,97 @@ +/** + * PyTorch Embedding Layer Node Definition + */ + +import { NodeDefinition } from '../../base' +import { NodeMetadata, BackendFramework } from '../../contracts' +import { TensorShape, BlockConfig, ConfigField, BlockType } from '../../../types' + +export class EmbeddingNode extends NodeDefinition { + readonly metadata: NodeMetadata = { + type: 'embedding', + label: 'Embedding', + category: 'advanced', + color: 'var(--color-purple)', + icon: 'TextAa', + description: 'Token embedding layer', + framework: BackendFramework.PyTorch + } + + readonly configSchema: ConfigField[] = [ + { + name: 'num_embeddings', + label: 'Vocabulary Size', + type: 'number', + required: true, + min: 1, + description: 'Size of the vocabulary' + }, + { + name: 'embedding_dim', + label: 'Embedding Dimension', + type: 'number', + required: true, + min: 1, + description: 'Size of each embedding vector' + }, + { + name: 'padding_idx', + label: 'Padding Index', + type: 'number', + default: -1, + description: 'Padding token index (or -1 for none)' + }, + { + name: 'max_norm', + label: 'Max Norm', + type: 'number', + default: 0, + min: 0, + description: 'Renormalize embeddings (0 for no normalization)' + }, + { + name: 'scale_grad_by_freq', + label: 'Scale Grad by Freq', + type: 'boolean', + default: false, + description: 'Scale gradients by word frequency' + } + ] + + computeOutputShape(inputShape: TensorShape | undefined, config: BlockConfig): TensorShape | undefined { + if (!inputShape || !config.embedding_dim) { + return undefined + } + + const embeddingDim = config.embedding_dim as number + + if (inputShape.dims.length === 2) { + const [batch, seqLen] = inputShape.dims as number[] + return { + dims: [batch, seqLen, embeddingDim], + description: `Embedding(${embeddingDim})` + } + } + + return { + dims: [...inputShape.dims, embeddingDim], + description: `Embedding(${embeddingDim})` + } + } + + validateIncomingConnection( + sourceNodeType: BlockType, + sourceOutputShape: TensorShape | undefined, + targetConfig: BlockConfig + ): string | undefined { + if (sourceNodeType === 'input' || sourceNodeType === 'dataloader') { + return undefined + } + + if (sourceNodeType === 'empty' || sourceNodeType === 'custom') { + return undefined + } + + return undefined + } +} diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/gru.ts b/project/frontend/src/lib/nodes/definitions/pytorch/gru.ts new file mode 100644 index 0000000..888c89d --- /dev/null +++ b/project/frontend/src/lib/nodes/definitions/pytorch/gru.ts @@ -0,0 +1,118 @@ +/** + * PyTorch GRU Layer Node Definition + */ + +import { NodeDefinition } from '../../base' +import { NodeMetadata, BackendFramework } from '../../contracts' +import { TensorShape, BlockConfig, ConfigField, BlockType } from '../../../types' + +export class GRUNode extends NodeDefinition { + readonly metadata: NodeMetadata = { + type: 'gru', + label: 'GRU', + category: 'advanced', + color: 'var(--color-purple)', + icon: 'ArrowsClockwise', + description: 'GRU recurrent layer', + framework: BackendFramework.PyTorch + } + + readonly configSchema: ConfigField[] = [ + { + name: 'hidden_size', + label: 'Hidden Size', + type: 'number', + required: true, + min: 1, + description: 'Number of features in hidden state' + }, + { + name: 'num_layers', + label: 'Layers', + type: 'number', + default: 1, + min: 1, + description: 'Number of recurrent layers' + }, + { + name: 'bias', + label: 'Use Bias', + type: 'boolean', + default: true, + description: 'Use bias weights' + }, + { + name: 'batch_first', + label: 'Batch First', + type: 'boolean', + default: true, + description: 'Input shape is (batch, seq, feature)' + }, + { + name: 'dropout', + label: 'Dropout', + type: 'number', + default: 0.0, + min: 0.0, + max: 1.0, + description: 'Dropout probability (if layers > 1)' + }, + { + name: 'bidirectional', + label: 'Bidirectional', + type: 'boolean', + default: false, + description: 'Use bidirectional GRU' + } + ] + + computeOutputShape(inputShape: TensorShape | undefined, config: BlockConfig): TensorShape | undefined { + if (!inputShape || !config.hidden_size) { + return undefined + } + + if (inputShape.dims.length !== 3) { + return undefined + } + + const batchFirst = config.batch_first ?? true + const hiddenSize = config.hidden_size as number + const bidirectional = config.bidirectional ?? false + + let batch: number, seqLen: number + + if (batchFirst) { + [batch, seqLen] = inputShape.dims as number[] + } else { + [seqLen, batch] = inputShape.dims as number[] + } + + const outFeatures = hiddenSize * (bidirectional ? 2 : 1) + + const dims = batchFirst ? [batch, seqLen, outFeatures] : [seqLen, batch, outFeatures] + + return { + dims, + description: `GRU(${hiddenSize})` + } + } + + validateIncomingConnection( + sourceNodeType: BlockType, + sourceOutputShape: TensorShape | undefined, + targetConfig: BlockConfig + ): string | undefined { + if (sourceNodeType === 'input' || sourceNodeType === 'dataloader') { + return undefined + } + + if (sourceNodeType === 'empty' || sourceNodeType === 'custom') { + return undefined + } + + return this.validateDimensions(sourceOutputShape, { + dims: 3, + description: '[batch, sequence, features] or [sequence, batch, features]' + }) + } +} diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/index.ts b/project/frontend/src/lib/nodes/definitions/pytorch/index.ts index 1afce9f..c6f70da 100644 --- a/project/frontend/src/lib/nodes/definitions/pytorch/index.ts +++ b/project/frontend/src/lib/nodes/definitions/pytorch/index.ts @@ -10,13 +10,20 @@ export { LossNode } from './loss' export { EmptyNode } from './empty' export { LinearNode } from './linear' export { Conv2DNode } from './conv2d' +export { Conv1DNode } from './conv1d' +export { Conv3DNode } from './conv3d' export { FlattenNode } from './flatten' export { ReLUNode } from './relu' export { DropoutNode } from './dropout' export { BatchNormNode } from './batchnorm' export { MaxPool2DNode } from './maxpool' +export { AvgPool2DNode } from './avgpool2d' +export { AdaptiveAvgPool2DNode } from './adaptiveavgpool2d' export { SoftmaxNode } from './softmax' export { ConcatNode } from './concat' export { AddNode } from './add' export { AttentionNode } from './attention' +export { LSTMNode } from './lstm' +export { GRUNode } from './gru' +export { EmbeddingNode } from './embedding' export { CustomNode } from './custom' diff --git a/project/frontend/src/lib/nodes/definitions/pytorch/lstm.ts b/project/frontend/src/lib/nodes/definitions/pytorch/lstm.ts new file mode 100644 index 0000000..12bdce7 --- /dev/null +++ b/project/frontend/src/lib/nodes/definitions/pytorch/lstm.ts @@ -0,0 +1,118 @@ +/** + * PyTorch LSTM Layer Node Definition + */ + +import { NodeDefinition } from '../../base' +import { NodeMetadata, BackendFramework } from '../../contracts' +import { TensorShape, BlockConfig, ConfigField, BlockType } from '../../../types' + +export class LSTMNode extends NodeDefinition { + readonly metadata: NodeMetadata = { + type: 'lstm', + label: 'LSTM', + category: 'advanced', + color: 'var(--color-purple)', + icon: 'ArrowsClockwise', + description: 'LSTM recurrent layer', + framework: BackendFramework.PyTorch + } + + readonly configSchema: ConfigField[] = [ + { + name: 'hidden_size', + label: 'Hidden Size', + type: 'number', + required: true, + min: 1, + description: 'Number of features in hidden state' + }, + { + name: 'num_layers', + label: 'Layers', + type: 'number', + default: 1, + min: 1, + description: 'Number of recurrent layers' + }, + { + name: 'bias', + label: 'Use Bias', + type: 'boolean', + default: true, + description: 'Use bias weights' + }, + { + name: 'batch_first', + label: 'Batch First', + type: 'boolean', + default: true, + description: 'Input shape is (batch, seq, feature)' + }, + { + name: 'dropout', + label: 'Dropout', + type: 'number', + default: 0.0, + min: 0.0, + max: 1.0, + description: 'Dropout probability (if layers > 1)' + }, + { + name: 'bidirectional', + label: 'Bidirectional', + type: 'boolean', + default: false, + description: 'Use bidirectional LSTM' + } + ] + + computeOutputShape(inputShape: TensorShape | undefined, config: BlockConfig): TensorShape | undefined { + if (!inputShape || !config.hidden_size) { + return undefined + } + + if (inputShape.dims.length !== 3) { + return undefined + } + + const batchFirst = config.batch_first ?? true + const hiddenSize = config.hidden_size as number + const bidirectional = config.bidirectional ?? false + + let batch: number, seqLen: number + + if (batchFirst) { + [batch, seqLen] = inputShape.dims as number[] + } else { + [seqLen, batch] = inputShape.dims as number[] + } + + const outFeatures = hiddenSize * (bidirectional ? 2 : 1) + + const dims = batchFirst ? [batch, seqLen, outFeatures] : [seqLen, batch, outFeatures] + + return { + dims, + description: `LSTM(${hiddenSize})` + } + } + + validateIncomingConnection( + sourceNodeType: BlockType, + sourceOutputShape: TensorShape | undefined, + targetConfig: BlockConfig + ): string | undefined { + if (sourceNodeType === 'input' || sourceNodeType === 'dataloader') { + return undefined + } + + if (sourceNodeType === 'empty' || sourceNodeType === 'custom') { + return undefined + } + + return this.validateDimensions(sourceOutputShape, { + dims: 3, + description: '[batch, sequence, features] or [sequence, batch, features]' + }) + } +}