| --- |
| 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). |
|
|
