Metadata-Version: 2.4
Name: neotune
Version: 0.9.1
Summary: Simple LLM fine-tuning with LoRA and DeepSpeed
License-Expression: MIT
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Intended Audience :: Science/Research
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch>=2.0
Requires-Dist: transformers>=4.40
Requires-Dist: peft>=0.10
Requires-Dist: datasets>=2.0
Requires-Dist: deepspeed>=0.12
Requires-Dist: accelerate>=0.30
Requires-Dist: scikit-learn>=1.0
Requires-Dist: numpy>=1.24
Requires-Dist: pyyaml>=6.0
Requires-Dist: python-dotenv>=1.0
Requires-Dist: tqdm>=4.60
Provides-Extra: ray
Requires-Dist: ray[train]>=2.0; extra == "ray"
Provides-Extra: logging
Requires-Dist: mlflow>=2.0; extra == "logging"
Requires-Dist: wandb>=0.15; extra == "logging"
Provides-Extra: all
Requires-Dist: neotune[logging,ray]; extra == "all"
Dynamic: license-file

# neotune

Simple LLM fine-tuning with LoRA and DeepSpeed.

Three inputs: a model, your datasets, and hyperparameters.

## Installation

```bash
pip install neotune
```

With optional extras:

```bash
pip install "neotune[ray]"       # Ray distributed training
pip install "neotune[logging]"   # MLflow + Weights & Biases
pip install "neotune[all]"       # everything
```

## Quick Start

```python
from neotune import finetune

results = finetune(
    model="google/gemma-3-4b-pt",
    datasets={"train": train_ds, "validation": val_ds},
    hyperparameters={"learning_rate": 2e-4, "num_train_epochs": 3},
)
```

Each dataset is a HuggingFace `Dataset` with a `"text"` column containing fully-formatted prompt/response strings, or pre-tokenized columns (`input_ids`, `attention_mask`, `labels`).

## Preparing Your Datasets

neotune expects you to bring your own HuggingFace `Dataset` objects.

### From a HuggingFace dataset

```python
from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-4b-pt")

ds = load_dataset("tatsu-lab/alpaca", split="train")

def format_example(example):
    messages = [
        {"role": "user", "content": example["instruction"]},
        {"role": "assistant", "content": example["output"]},
    ]
    return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}

ds = ds.map(format_example, remove_columns=ds.column_names)
splits = ds.train_test_split(test_size=0.1, seed=42)

# splits["train"] and splits["test"] each have a "text" column
```

### From a CSV file

```python
from datasets import load_dataset
from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("google/gemma-3-4b-pt")

ds = load_dataset("csv", data_files="data.csv", split="train")

def format_example(example):
    messages = [
        {"role": "user", "content": example["prompt"]},
        {"role": "assistant", "content": example["response"]},
    ]
    return {"text": tokenizer.apply_chat_template(messages, tokenize=False)}

ds = ds.map(format_example, remove_columns=ds.column_names)
splits = ds.train_test_split(test_size=0.1, seed=42)
```

## API Reference

### `finetune(model, datasets, hyperparameters)` -> `dict`

One-call convenience function. Returns test-set metrics if a `"test"` split was provided, otherwise an empty dict.

```python
from neotune import finetune

results = finetune(
    model="google/gemma-3-4b-pt",
    datasets={"train": train_ds, "validation": val_ds, "test": test_ds},
    hyperparameters={"learning_rate": 2e-4},
)
```

### `NeoTune(model, datasets, hyperparameters)`

Class-based API.

```python
from neotune import NeoTune

nt = NeoTune(
    model="google/gemma-3-4b-pt",
    datasets={"train": train_ds, "validation": val_ds},
    hyperparameters={"num_train_epochs": 5, "output_dir": "./my-adapter"},
)

results = nt.train()
```

#### Parameters

**`model`** -- `str`
HuggingFace model ID or local path.

**`datasets`** -- `dict[str, Dataset]`
A dict of HuggingFace `Dataset` objects. `"train"` is required. `"validation"` and `"test"` are optional.

**`hyperparameters`** -- `dict`, optional
Override any default. All keys are optional:

| Key | Default | Description |
|-----|---------|-------------|
| **Training** | | |
| `learning_rate` | `1e-4` | Learning rate |
| `num_train_epochs` | `3` | Number of training epochs |
| `batch_size` | `1` | Per-device batch size |
| `gradient_accumulation_steps` | `4` | Gradient accumulation steps |
| `warmup_ratio` | `0.03` | Warmup ratio |
| `weight_decay` | `0.01` | Weight decay |
| `bf16` | `True` | bfloat16 mixed precision |
| `gradient_checkpointing` | `False` | Gradient checkpointing |
| `logging_steps` | `10` | Log every N steps |
| `eval_steps` | `50` | Evaluate every N steps |
| `save_steps` | `100` | Checkpoint every N steps |
| `save_total_limit` | `3` | Max checkpoints to keep |
| **LoRA** | | |
| `lora_r` | `16` | LoRA rank |
| `lora_alpha` | `32` | LoRA alpha |
| `lora_dropout` | `0.05` | LoRA dropout |
| `lora_target_modules` | `"all-linear"` | Target modules (auto-detects all linear layers) |
| **Output** | | |
| `output_dir` | `"./adapter-output"` | Where to save the adapter |
| `hf_repo` | `None` | Push to HuggingFace Hub |
| **DeepSpeed** | | |
| `ds_config` | `None` | `None` (auto: DeepSpeed ZeRO-2 on multi-GPU, disabled on single GPU), `"auto"` (force ZeRO-2), `False` (force off), a file path, or a dict |
| **Data** | | |
| `max_len` | `2048` | Max sequence length |
| `dataset_text_field` | `"text"` | Column name for training text |

#### Methods

- **`.train()`** -> `dict` -- Fine-tune and return test metrics (if test split provided).
- **`.tokenizer`** -- Access the underlying tokenizer.

## Device & DeepSpeed

neotune auto-detects your hardware:

| Device | Default behavior |
|--------|-----------------|
| CPU | Trains on CPU (float32) |
| Single CUDA GPU | Standard GPU training (bf16) |
| Multi CUDA GPU | DeepSpeed ZeRO-2 auto-enabled (bf16) |
| Apple Silicon (MPS) | MPS acceleration (fp16, no DeepSpeed) |

Override with `ds_config`:

```python
# Force DeepSpeed off (e.g. multi-GPU notebook without mpi4py)
finetune(model, datasets, {"ds_config": False})

# Force DeepSpeed on (even on single GPU)
finetune(model, datasets, {"ds_config": "auto"})

# Custom DeepSpeed config
finetune(model, datasets, {"ds_config": "my_ds_config.json"})
```

> **Notebook users:** If you have multiple GPUs but get `ModuleNotFoundError: No module named 'mpi4py'`, either install it (`pip install mpi4py`) or disable DeepSpeed with `"ds_config": False`.

## Advanced Usage

### Generative evaluation

```python
from neotune.eval import generate_and_evaluate

results = generate_and_evaluate(
    model_id="google/gemma-3-4b-pt",
    adapter_dir="./my-adapter",
    test_ds=test_ds,
    prompt_col="instruction",
    label_col="expected_output",
)
```

### Distributed training with DeepSpeed (CLI)

```bash
deepspeed --num_gpus 4 -m neotune.train --config config.yaml --mode train
```

### Distributed training with Ray

```bash
python -m neotune.ray_train --config config.yaml --num_workers 4
```

### Kubernetes (KubeRay)

See `k8s/rayjob-lora-sft.yaml` for a KubeRay `RayJob` template.

## Environment Variables

| Variable | Description |
|----------|-------------|
| `HF_TOKEN` | HuggingFace access token (for gated models) |
| `WANDB_API_KEY` | Weights & Biases API key (optional) |

Create a `.env` file in your working directory:

```
HF_TOKEN=your_token_here
WANDB_API_KEY=your_wandb_key_here
```

## License

MIT
