| """Gradio web frontend for MiniRewrite. |
| Self-contained application ready for Hugging Face Spaces. |
| """ |
|
|
| import sys |
| import time |
| from pathlib import Path |
| import gradio as gr |
| import torch |
|
|
| |
| ROOT_DIR = Path(__file__).resolve().parent |
| if str(ROOT_DIR / "src") not in sys.path: |
| sys.path.insert(0, str(ROOT_DIR / "src")) |
|
|
| from minirewrite.pipeline import MiniRewrite |
|
|
| |
| print("Initializing MiniRewrite model...") |
| MODEL_DIR = ROOT_DIR if (ROOT_DIR / "best.pt").exists() else (ROOT_DIR / "runs" / "rtx4060") |
|
|
| |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| pipeline = MiniRewrite.from_pretrained(MODEL_DIR, checkpoint_name="best.pt", device=device) |
| print(f"Model loaded successfully on {pipeline.device}!") |
|
|
|
|
| def process_rewrite(text: str, task: str, tone: str, style: str): |
| if not text or not text.strip(): |
| return "Please enter a sentence to rewrite.", "0 ms" |
|
|
| start_time = time.perf_counter() |
| result = pipeline.rewrite( |
| text.strip(), |
| task=task.lower(), |
| tone=tone.lower(), |
| style=style.lower(), |
| max_source_length=128, |
| max_new_tokens=128, |
| ) |
| elapsed_ms = (time.perf_counter() - start_time) * 1000 |
|
|
| latency_str = f"Generated in {elapsed_ms:.1f} ms on {pipeline.device}" |
| return result, latency_str |
|
|
|
|
| theme = gr.themes.Soft( |
| primary_hue="indigo", |
| secondary_hue="purple", |
| ) |
|
|
| with gr.Blocks(theme=theme, title="MiniRewrite - Transformer Text Rewriter") as demo: |
| with gr.Column(elem_id="title-banner"): |
| gr.Markdown( |
| """ |
| # ✍️ MiniRewrite |
| **Instruction-Controlled Encoder-Decoder Transformer (7.6M Parameters)** |
| |
| Rewrites text by controlling **Task**, **Tone**, and **Style** in real time. |
| """ |
| ) |
|
|
| with gr.Row(): |
| with gr.Column(scale=5): |
| input_text = gr.Textbox( |
| label="Input Sentence", |
| placeholder="e.g. send me report asap", |
| lines=3, |
| max_lines=5, |
| ) |
|
|
| with gr.Row(): |
| task_dropdown = gr.Dropdown( |
| label="Task", |
| choices=["rewrite", "simplify", "shorten"], |
| value="rewrite", |
| ) |
| tone_dropdown = gr.Dropdown( |
| label="Tone", |
| choices=["professional", "friendly", "neutral", "casual"], |
| value="professional", |
| ) |
| style_dropdown = gr.Dropdown( |
| label="Style", |
| choices=["polite", "direct", "concise", "plain"], |
| value="polite", |
| ) |
|
|
| submit_btn = gr.Button("Rewrite Sentence", variant="primary", size="lg") |
|
|
| with gr.Column(scale=5): |
| output_text = gr.Textbox( |
| label="Rewritten Output", |
| lines=3, |
| max_lines=5, |
| interactive=False, |
| ) |
| latency_badge = gr.Markdown("Latency: ready") |
|
|
| submit_btn.click( |
| fn=process_rewrite, |
| inputs=[input_text, task_dropdown, tone_dropdown, style_dropdown], |
| outputs=[output_text, latency_badge], |
| ) |
|
|
| gr.Markdown("### 💡 Click any preset example to test instantly:") |
| gr.Examples( |
| examples=[ |
| ["send me report asap", "rewrite", "professional", "polite"], |
| ["send me report asap", "rewrite", "friendly", "direct"], |
| ["cant come to meeting tmrw", "rewrite", "professional", "polite"], |
| ["pls fix the bug today", "rewrite", "friendly", "direct"], |
| ["In order to complete the project on time, we need more resources.", "shorten", "neutral", "concise"], |
| ["There was a concert of Michael Jackson in my city last month .", "simplify", "neutral", "plain"], |
| ], |
| inputs=[input_text, task_dropdown, tone_dropdown, style_dropdown], |
| outputs=[output_text, latency_badge], |
| fn=process_rewrite, |
| cache_examples=False, |
| ) |
|
|
| if __name__ == "__main__": |
| demo.launch() |
|
|
