A collection of tools for machine learning research, including experiment management, Kubernetes utilities, and LaTeX processing.
-
LaTeX Tools
- Grammar and style checker for LaTeX documents
- Automatic suggestions for improving technical writing
-
LLM Integration
- Easy interaction with OpenAI API compatible LLMs
- Support for multiple model presets and tiers
-
Document Tools
- Ask questions and get answers about documentation
-
Experiment Management
- Weights & Biases run logs downloader
-
Kubernetes Tools
- Pod port forwarding with automatic reconnection
-
Caching System
- Redis-based function result caching
- Transparent caching with decorators
pip install ml_research_toolsgit clone https://github.com/alexdremov/ml_research_tools.git
cd ml_research_tools
poetry installThe toolkit can be configured through multiple methods, with a cascading priority:
- Command-line arguments (highest priority)
- Configuration file
- Default values (lowest priority)
By default, the configuration is stored in ~/.config/ml_research_tools/config.yaml.
If this file doesn't exist, it will be created with default values when the tool is first run.
Example configuration file:
logging:
level: INFO
redis:
host: localhost
port: 6379
db: 0
enabled: true
ttl: 604800 # 7 days in seconds
llm:
default: "openai" # Default preset to use
presets:
openai:
base_url: https://api.openai.com/v1
model: gpt-3.5-turbo
max_tokens: 8000
temperature: 0.01
top_p: 1.0
retry_attempts: 3
retry_delay: 5
api_key: null
tier: standard
ollama:
model: gemma3
base_url: http://127.0.0.1:3333/v1/
perplexity:
base_url: https://api.perplexity.ai/
model: sonar-pro
max_tokens: 128000
temperature: 0.01
api_key: null
tier: premiumConfiguration can be overridden using command-line arguments:
ml_research_tools --log-level DEBUG --redis-host redis.example.com --llm-preset premium paper.tex latex-grammar paper.texUsage: ml_research_tools [-h] [--config CONFIG]
[--log-level {DEBUG,INFO,WARNING,ERROR,CRITICAL}]
[--log-file LOG_FILE] [-v] [--redis-host REDIS_HOST]
[--redis-port REDIS_PORT] [--redis-db REDIS_DB]
[--redis-password REDIS_PASSWORD] [--redis-disable]
[--redis-recache] [--llm-preset LLM_PRESET]
[--llm-tier LLM_TIER] [--llm-api-key LLM_API_KEY]
[--llm-base-url LLM_BASE_URL] [--llm-model LLM_MODEL]
[--llm-max-tokens LLM_MAX_TOKENS]
[--llm-temperature LLM_TEMPERATURE]
[--llm-top-p LLM_TOP_P]
[--llm-retry-attempts LLM_RETRY_ATTEMPTS]
[--llm-retry-delay LLM_RETRY_DELAY] [--list-presets]
[--list-tools]
{help,ask-document,wandb-downloader,kube-pod-forward,latex-grammar}
...
ML Research Tools - A collection of utilities for ML research
Options:
-h, --help show this help message and exit
--list-presets List available LLM presets and exit
--list-tools List available tools and exit
Configuration:
--config CONFIG Path to configuration file (default:
~/.config/ml_research_tools/config.yaml)
Logging:
--log-level {DEBUG,INFO,WARNING,ERROR,CRITICAL}
Logging level
--log-file LOG_FILE Path to log file
-v, --verbose Enable verbose logging
Redis:
--redis-host REDIS_HOST
Redis host
--redis-port REDIS_PORT
Redis port
--redis-db REDIS_DB Redis database number
--redis-password REDIS_PASSWORD
Redis password
--redis-disable Disable Redis caching
--redis-recache Disable Redis caching retrieval, but allow saving
Llm:
--llm-preset LLM_PRESET
LLM preset name to use (e.g., 'standard', 'premium')
--llm-tier LLM_TIER LLM tier to use (e.g., 'standard', 'premium')
--llm-api-key LLM_API_KEY
API key for LLM service
--llm-base-url LLM_BASE_URL
Base URL for the LLM API endpoint
--llm-model LLM_MODEL
LLM model to use
--llm-max-tokens LLM_MAX_TOKENS
Maximum tokens for LLM response
--llm-temperature LLM_TEMPERATURE
Temperature for LLM sampling
--llm-top-p LLM_TOP_P
Top-p value for LLM sampling
--llm-retry-attempts LLM_RETRY_ATTEMPTS
Number of retry attempts for LLM API calls
--llm-retry-delay LLM_RETRY_DELAY
Delay between retry attempts for LLM API calls
(seconds)
- Clone the repository
- Install development dependencies:
poetry install --with dev
- Run tests:
poetry run pytest
- Code quality tools:
# Format code poetry run black . poetry run isort . # Check typing poetry run mypy . # Run linter poetry run ruff .
- Create a new module in the appropriate directory
- Implement a class that inherits from
BaseTool - Register arguments in the
add_argumentsmethod - Implement the
executemethod - Import the module in
__init__.pyto ensure discovery
Example:
from ml_research_tools.core.base_tool import BaseTool
class MyNewTool(BaseTool):
name = "my-tool"
description = "Description of my new tool"
@classmethod
def add_arguments(cls, parser):
parser.add_argument("--option", help="An option for my tool")
def execute(self, config, args):
# Implementation
return 0 # SuccessThis project is licensed under the MIT License - see the LICENSE file for details.
The full documentation is available at https://alexdremov.github.io/ml_research_tools/.
To build the documentation locally:
poetry install --with docs
cd docs
poetry run make htmlThen open docs/build/html/index.html in your browser.
This project is wildly LLM-written (though, widely reviewed by me)