Metadata-Version: 2.4
Name: bantam-cli
Version: 0.1.3
Summary: Bantam language model core
Author-email: Theoistic <theo@theoistic.com>
Requires-Python: >=3.8
Description-Content-Type: text/markdown
Requires-Dist: torch>=2.1.0
Requires-Dist: transformers>=4.40.0
Requires-Dist: datasets>=2.16.0
Requires-Dist: peft>=0.11.0
Requires-Dist: numpy>=1.24.0
Requires-Dist: matplotlib>=3.8.0
Requires-Dist: tqdm>=4.65.0
Requires-Dist: PyYAML>=6.0
Requires-Dist: typer>=0.9.0

# Bantam

Bantam is a configurable decoder-only transformer packaged with a CLI for tokenizer training, pre-training, SFT, and chat inference. It runs on pure PyTorch (`torch>=2.1`) and uses the scaled-dot-product/FlashAttention stack for grouped-query attention, sliding windows, and sink tokens without custom kernels.

## Key Features
- Flex-attention based GQA with optional sliding windows and sink tokens.
- Per-layer overrides let you mix head counts, windows, sinks, and MoE routing per block via `layer_configs`.
- SwiGLU dense blocks or Switch/Top-k MoE with router/aux loss controls.
- Rotary embeddings with linear, NTK, and YaRN-style scaling.
- Regex-aware BPE trainer (`bantam-cli trainbpe`) that isolates punctuation, chunks long numbers, and optionally prefills 0–N.
- Unified CLI for model init, pre-training, SFT (full or LoRA), resume, and chat streaming.

## Installation

```bash
python3 -m pip install bantam-cli
```

## CLI Quick Start

```bash
bantam-cli --help          # list commands

# Train a tokenizer from JSONL
bantam-cli trainbpe --input data/pretrain.jsonl --out tokenizers/run1

# Generate starter configs (model + train + SFT)
bantam-cli define --out configs

# Initialise a model directory from split configs
bantam-cli init --mconf configs/models/pretrain_small.yaml --config configs/training/pretrain_small.yaml --out models/bantam-init

# Launch pre-training run
bantam-cli pretrain --mconf configs/models/pretrain_small.yaml --config configs/training/pretrain_small.yaml

# Resume a checkpointed run (model config auto-detected)
bantam-cli continue-pretrain --config configs/training/pretrain_small.yaml --model models/run1/checkpoints/step_20000

# Supervised fine-tuning (LoRA or full)
bantam-cli sft --config configs/training/sft_small.yaml --model models/bantam-init

# Chat with a trained checkpoint (pretrain format by default)
bantam-cli chat --model models/run1/best

# base model inference
bantam-cli chat --model Theoistic/Bantam-285m

# Conversational chat retaining the last two turns
bantam-cli chat --model Theoistic/Bantam-285m-Instruct --format chat --history-turns 2
```

Model configs live under `configs/models/` with a `bantam_config` block. Training/SFT configs live under `configs/training/` with `training_args`, optional `sft_args`, and a run `seed`. Relative paths resolve against the YAML file location.

Example per-layer tweak:

```yaml
bantam_config:
  hidden_size: 1024
  num_hidden_layers: 12
  num_attention_heads: 16
  num_key_value_heads: 4
  layer_configs:
    - { window: 2048, num_attention_sinks: 2 }
    - { expert_type: switch, num_experts: 8, moe_capacity_factor: 1.0 }
```

## More Info

- Source code & full documentation: https://github.com/theo-bntm/bantam
- Release checklist: `publish.md` inside the repository

Report issues or ideas via GitHub. Enjoy experimenting with Bantam!

## LLaMAFactory Integration

You can fine-tune Bantam models with [LLaMAFactory](https://github.com/hiyouga/LLaMA-Factory) once `bantam-cli` is installed 

you need to register the architecture in LLaMA-Factory/src/llamafactory/cli.py

```python
import bantam  # lazy imports
import bantam.tokenization_bantam  # registers BantamFastTokenizer with AutoTokenizer
import bantam.modeling_bantam      # registers config/model with AutoConfig/AutoModel
```

and add this somewhere in the bottom of LLaMA-Factory/src/llamafactory/data/template.py

```python
register_template(
    name="bantam_chat",
    format_user=StringFormatter(
        slots=["<|USER_START|>{{content}}<|USER_END|><|AGENT_START|>"]
    ),
    format_assistant=StringFormatter(slots=["{{content}}<|AGENT_END|>"]),
    format_system=StringFormatter(slots=["<|SYSTEM_START|>{{content}}<|SYSTEM_END|>"]),
    format_prefix=EmptyFormatter(slots=[{"bos_token"}]),
    stop_words=["<|AGENT_END|>", "<|EOS|>"],
    replace_eos=False,
)
```

- Use `--template bantam_chat` in LLaMAFactory for both SFT and preference/DPO runs.
- The template leaves user turns in-place but only supervises the assistant span (mirrors Bantam’s `mask_user_queries=true` default).
- Set `include_agent_end=true` and `include_eos=true` in your Bantam SFT config so completions end with `<|AGENT_END|><|EOS|>`, giving the stop tokens above reliable cut points.
- For LoRA vs. full SFT, follow LLaMAFactory’s standard CLI flags—Bantam checkpoints expose all weights so “full” fine-tunes work out of the box.

After training, export the adapter or merged weights from LLaMAFactory and load them with `bantam-cli chat --model <path>` to validate generations.
