YAML Metadata Warning:The pipeline tag "text2text-generation" is not in the official list: text-classification, token-classification, table-question-answering, question-answering, zero-shot-classification, translation, summarization, feature-extraction, text-generation, fill-mask, sentence-similarity, text-to-speech, text-to-audio, automatic-speech-recognition, audio-to-audio, audio-classification, audio-text-to-text, voice-activity-detection, depth-estimation, image-classification, object-detection, image-segmentation, text-to-image, image-to-text, image-to-image, image-to-video, unconditional-image-generation, video-classification, reinforcement-learning, robotics, tabular-classification, tabular-regression, tabular-to-text, table-to-text, multiple-choice, text-ranking, text-retrieval, time-series-forecasting, text-to-video, image-text-to-text, image-text-to-image, image-text-to-video, visual-question-answering, document-question-answering, zero-shot-image-classification, graph-ml, mask-generation, zero-shot-object-detection, text-to-3d, image-to-3d, image-feature-extraction, video-text-to-text, keypoint-detection, visual-document-retrieval, any-to-any, video-to-video, other

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

Downloads last month: -; Downloads are not tracked for this model. How to track