File size: 4,247 Bytes
7f90954 1bedb07 7f90954 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """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
# Ensure local src package is on path
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
# Global model instance
print("Initializing MiniRewrite model...")
MODEL_DIR = ROOT_DIR if (ROOT_DIR / "best.pt").exists() else (ROOT_DIR / "runs" / "rtx4060")
# CPU inference is extremely fast (~50ms) for this 7.6M model and requires zero GPU cold-start
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()
|