metadata
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
pip install torch tokenizers pydantic pyyaml gradio
2. Download & Run with 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
python app.py
4. Interactive CLI
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).