gnn-ml-integration · git:20260906.e62b9e9 · 2026-09-06 · sha256 4e6082785ded46d9
gnn-ml-integration git:20260906.e62b9e9A
Immutable. This exact content is served forever at /api/v1/blob/4e6082785ded46d9.
---
name: gnn-ml-integration
description: GNN machine learning integration and model training. Use when training ML models on GNN data, checking ML framework availability, or integrating GNN pipeline outputs with machine learning workflows.
---
# GNN Machine Learning Integration (Step 14)
## Purpose
Integrates machine learning capabilities with the GNN pipeline, enabling feature extraction from parsed models, ML model training on GNN data, and framework availability checking.
## Key Commands
```bash
# Run ML integration
python src/gnn/14_ml_integration.py --target-dir input/gnn_files --output-dir output --verbose
# As part of pipeline
python src/gnn/main.py --only-steps 14 --verbose
```
## API
```python
from gnn.ml_integration import process_ml_integration, check_ml_frameworks, feature_vector, predict_with_model
# Process ML integration step (used by pipeline)
result = process_ml_integration(target_dir, output_dir, verbose=True)
# Check available ML frameworks
frameworks = check_ml_frameworks()
# Returns: {'pytorch': {'available': True, 'version': '2.x'},
# 'tensorflow': {'available': False},
# 'jax': {'available': True, 'version': '0.4.x'},
# 'sklearn': {'available': True, 'version': '1.x'}}
# Build the canonical feature vector and predict from a saved .pkl artifact
vec = feature_vector(features)
label = predict_with_model(Path("output/14_ml_integration_output/gnn_decision_tree.pkl"), features, label_names=labels)
```
## Key Exports
- `process_ml_integration` — main pipeline processing function
- `check_ml_frameworks` — checks availability of PyTorch, TensorFlow, JAX, scikit-learn
- `extract_gnn_features` — structural feature extraction for one GNN file
- `feature_vector` — canonical numeric vector in `NUMERIC_FEATURE_NAMES` order
- `complexity_label` — small/medium/large via `COMPLEXITY_THRESHOLDS` (100, 1000)
- `summarize_features` — min/max/mean over `SUMMARY_STATISTIC_KEYS`
- `load_classifier`, `predict_with_model`, `predict_batch` — inference from saved `.pkl` artifacts (scikit-learn required only at call time; `label_names` comes from `ml_integration_results.json`)
- `InferenceError` — raised for missing/invalid artifacts or prediction failures
- `get_module_info` — version, feature flags, tool inventory
- Constants: `NUMERIC_FEATURE_NAMES`, `COMPLEXITY_THRESHOLDS`, `COMPLEXITY_LABELS`, `SUMMARY_STATISTIC_KEYS`
## Supported ML Frameworks
| Framework | Check Key | Purpose |
| ----------- | ----------- | --------- |
| **PyTorch** | `pytorch` | Deep learning, CUDA support detection |
| **TensorFlow** | `tensorflow` | Neural network training |
| **JAX** | `jax` | Differentiable computing |
| **scikit-learn** | `sklearn` | Classical ML algorithms |
## Dependencies
```bash
# ML training deps
uv sync --extra ml-ai
# Includes: transformers, scipy, scikit-learn
# Note: torch ships in its own `torch` extra (uv sync --extra torch;
# torch>=2.13.0 resolves GHSA-rrmf-rvhw-rf47).
```
## Output
- ML results in `output/14_ml_integration_output/`
- Framework availability reports
## MCP Tools
This module registers tools with the GNN MCP server (see `mcp.py`):
- `check_ml_frameworks`
- `get_ml_module_info`
- `list_ml_integration_targets`
- `process_ml_integration`
## References
- [AGENTS.md](AGENTS.md) — Module documentation
- [README.md](README.md) — Usage guide
- [SPEC.md](SPEC.md) — Module specification
---
## Documentation
- **[README](README.md)**: Module Overview
- **[AGENTS](AGENTS.md)**: Agentic Workflows
- **[SPEC](SPEC.md)**: Architectural Specification
- **[SKILL](SKILL.md)**: Capability API