---
license: mit
tags:
- pytorch
- text2text-generation
- style-transfer
- transformer
- seq2seq
pipeline_tag: text2text-generation
language:
- en
---

# ✍️ MiniRewrite — Instruction-Controlled Transformer (7.6M)

MiniRewrite is a custom **7.6M-parameter Encoder-Decoder Transformer** built from scratch in PyTorch for controlled text rewriting, tone transfer, and text simplification.

- **Fast & Lightweight:** ~15 ms latency on GPU, ~70 ms on CPU (~91 MB checkpoint).
- **Controllable Attributes:** Task, Tone, and Style are conditioned via special prefix instruction tokens.

---

## 🎛️ Supported Controls

| Dimension | Available Options |
| :--- | :--- |
| **Task** | `rewrite`, `simplify`, `shorten` |
| **Tone** | `professional`, `friendly`, `neutral`, `casual` |
| **Style** | `polite`, `direct`, `concise`, `plain` |

---

## 🚀 Quickstart

### 1. Installation
```bash
pip install torch tokenizers pydantic pyyaml gradio
```

### 2. Download & Run with Python
```python
from huggingface_hub import snapshot_download
import sys
from pathlib import Path

# Download model from Hugging Face
repo_path = snapshot_download(repo_id="siddheshrj/minirewrite")
sys.path.insert(0, str(Path(repo_path) / "src"))

from minirewrite.pipeline import MiniRewrite

# Load model in 1 line
model = MiniRewrite.from_pretrained(repo_path, checkpoint_name="best.pt", device="auto")

# Rewrite with control
output = model.rewrite(
    "send me report asap",
    task="rewrite",
    tone="professional",
    style="polite"
)
print("Rewritten:", output)
# Output: Could you please send me report asap?
```

### 3. Run the Gradio Frontend Web UI
```bash
python app.py
```

### 4. Interactive CLI
```bash
python predict.py "send me report asap" --tone professional --style polite
python predict.py --interactive
```

---

## 🧠 Architecture & Training
- **Encoder-Decoder Transformer:** 6 encoder layers, 6 decoder layers, 8 attention heads, $d_{\text{model}}=256$, $d_{\text{ff}}=1024$.
- **Tokenizer:** Custom Byte-Level BPE tokenizer with 8,000 vocabulary size and dedicated instruction control tokens.
- **Trained on:** 61,074 high-quality instruction-aligned rewrite pairs.
- **Hardware:** Trained locally on NVIDIA GeForce RTX 4060 Laptop GPU.
- **Validation Loss:** 0.9410 (Perplexity ~2.56).