ml-research

ML Research Plugin for Claude Code

Comprehensive Claude Code plugin for machine learning research and experimentation. Provides project scaffolding, experiment management, training support, and debugging tools for PyTorch Lightning, Hydra, PyTorch Geometric, and Hugging Face Transformers.

Installation

Via Claude Code Marketplace

# Add marketplace
/plugin marketplace add nishide-dev/claude-code-ml-research

# Install plugin
/plugin install ml-research

Via GitHub URL

/plugin install gh:nishide-dev/claude-code-ml-research

Via Local Clone

git clone https://github.com/nishide-dev/claude-code-ml-research.git
/plugin install ./claude-code-ml-research

Features

Skills

This plugin provides comprehensive ML workflow skills following the Agent Skills open standard. All skills are directory-based with supporting files (templates, examples, scripts).

Workflow Skills (user-invocable commands):

• /ml-train - Execute training runs with proper monitoring, checkpointing, and experiment tracking
• /ml-config-manager - Generate and manage Hydra configuration files (model, data, trainer, logger, experiment configs)
• /ml-debug - Debug common ML training issues (NaN loss, OOM, slow training, convergence problems)
• /ml-experiment - Manage ML experiments, track results, and compare performance across configurations
• /ml-validate - Comprehensive validation of ML project structure, configurations, and training readiness
• /ml-profile - Profile ML training performance to identify bottlenecks and optimize GPU utilization
• /ml-data-pipeline - Create and manage data loading, preprocessing, and augmentation pipelines
• /ml-setup - Setup development environment with modern Python tooling (uv/pixi)
• /ml-project-init - Initialize new ML research project using the ML Research template
• /ml-lint - Run comprehensive code quality checks (ruff format/check, type checking)
• /ml-format - Format code with ruff and verify results
• /ml-model-export - Export trained PyTorch models to various formats (ONNX, TorchScript, TensorRT)

Knowledge Skills (background knowledge loaded automatically):

• /ml-lightning-basics - PyTorch Lightning patterns and best practices
• /ml-hydra-config - Hydra configuration management
• /ml-pytorch-geometric - Graph Neural Networks with PyTorch Geometric
• /ml-wandb-tracking - Experiment tracking with Weights & Biases
• /ml-transformers - Hugging Face Transformers + Lightning integration
• /ml-cli-tools - Building CLIs with Typer and Rich
• /tool-pixi - Pixi package manager for ML projects
• /tool-marimo - Marimo reactive Python notebooks

Agents

Specialized sub-agents with visual color coding for easy identification:

• ml-architect 🔵 (Opus): Design ML system architectures and model designs
• training-debugger 🔴 (Sonnet): Diagnose and fix training issues
• config-generator 🟣 (Sonnet): Generate and validate Hydra configurations
• pytorch-expert 🟠 (Sonnet): Write efficient PyTorch code and optimize performance
• geometric-specialist 🔷 (Sonnet): Implement Graph Neural Networks with PyTorch Geometric
• transformers-specialist 💜 (Sonnet): Hugging Face Transformers expert for LLM fine-tuning, PEFT (LoRA/QLoRA), tokenization, and NLP tasks

Rules

Automatically enforced best practices for ML development:

• coding-standards: Deterministic operations, type hints, tensor shape documentation, PyTorch Lightning patterns
• security-practices: API key protection, environment variables, file path sanitization, sensitive data handling
• workflow-constraints: Config validation, checkpointing requirements, experiment tracking, reproducibility

Hooks

Event-driven automation for code quality:

• Auto-format Python files with Ruff after Write/Edit
• Validate YAML configs after writing Hydra configs
• Remind dependencies installation after package config changes
• Welcome message on session start with available commands

Supported Frameworks

• PyTorch Lightning: High-level training framework
• PyTorch: Core deep learning framework
• PyTorch Geometric: Graph Neural Networks
• Hugging Face Transformers: NLP and vision transformers
• Hydra: Configuration management
• Weights & Biases: Experiment tracking

Supported Tools

• pixi: Conda-based package manager with PyPI integration (recommended for GPU/CUDA projects)
• uv: Fast Python package manager with pip compatibility (recommended for CPU or simple projects)
• ruff: Linting and formatting
• ty: Type checking

Development

Local Development

# Clone repository
git clone https://github.com/nishide-dev/claude-code-ml-research.git
cd claude-code-ml-research

# Install ML-specific rules (recommended for all ML projects)
cp -r rules/ml/* ~/.claude/rules/

# Test plugin locally
claude --plugin-dir . code