| """Standalone prediction script for MiniRewrite. |
| Works on any machine with PyTorch and Tokenizers installed (CPU or CUDA). |
| |
| Usage: |
| # 1. Single prediction: |
| python predict.py "send me report asap" --tone professional --style polite |
| |
| # 2. Friendly direct: |
| python predict.py "send me report asap" --tone friendly --style direct |
| |
| # 3. Shorten concise: |
| python predict.py "In order to complete the project on time, we need more resources." --task shorten --style concise |
| |
| # 4. Interactive chat mode: |
| python predict.py --interactive |
| """ |
|
|
| import argparse |
| import sys |
| from pathlib import Path |
| import torch |
|
|
| |
| SCRIPT_DIR = Path(__file__).resolve().parent |
| if str(SCRIPT_DIR / "src") not in sys.path: |
| sys.path.insert(0, str(SCRIPT_DIR / "src")) |
|
|
| from minirewrite.tokenizer.tokenizer import RewriteTokenizer |
| from minirewrite.model.transformer import Seq2SeqTransformer, ModelConfig |
| from minirewrite.inference.generate import rewrite |
| from minirewrite.training.trainer import load_model_for_inference |
|
|
|
|
| def load_model(run_dir: Path, checkpoint_name: str, device: torch.device): |
| print(f"Loading model and tokenizer from {run_dir} ({checkpoint_name}) to {device}...") |
| return load_model_for_inference(run_dir, device, checkpoint_name) |
|
|
|
|
| def run_interactive(model, tokenizer, device): |
| print("\n" + "=" * 60) |
| print("MiniRewrite Interactive Mode (type 'exit' or 'quit' to stop)") |
| print("=" * 60) |
|
|
| tasks = ["rewrite", "simplify", "shorten"] |
| tones = ["professional", "friendly", "neutral", "casual"] |
| styles = ["polite", "direct", "concise", "plain"] |
|
|
| current_task = "rewrite" |
| current_tone = "professional" |
| current_style = "polite" |
|
|
| print(f"Current defaults: task={current_task}, tone={current_tone}, style={current_style}") |
| print("Type '!settings' to change task/tone/style.\n") |
|
|
| while True: |
| try: |
| text = input("Input text > ").strip() |
| except (KeyboardInterrupt, EOFError): |
| break |
|
|
| if not text: |
| continue |
| if text.lower() in ("exit", "quit", "q"): |
| print("Goodbye!") |
| break |
| if text == "!settings": |
| print(f"Tasks: {', '.join(tasks)}") |
| t = input(f"Select task [{current_task}]: ").strip() |
| if t in tasks: |
| current_task = t |
|
|
| print(f"Tones: {', '.join(tones)}") |
| to = input(f"Select tone [{current_tone}]: ").strip() |
| if to in tones: |
| current_tone = to |
|
|
| print(f"Styles: {', '.join(styles)}") |
| st = input(f"Select style [{current_style}]: ").strip() |
| if st in styles: |
| current_style = st |
| print(f"Updated: task={current_task}, tone={current_tone}, style={current_style}\n") |
| continue |
|
|
| req = [{"source": text, "task": current_task, "tone": current_tone, "style": current_style}] |
| out = rewrite(model, tokenizer, req, max_source_length=128, max_new_tokens=128)[0] |
| print(f"Rewritten > {out}\n") |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description="Test MiniRewrite on any text.") |
| parser.add_argument("text", nargs="?", help="Text to rewrite") |
| parser.add_argument("--interactive", "-i", action="store_true", help="Launch interactive prompt") |
| parser.add_argument("--task", default="rewrite", choices=["rewrite", "simplify", "shorten"]) |
| parser.add_argument("--tone", default="professional", choices=["professional", "friendly", "neutral", "casual"]) |
| parser.add_argument("--style", default="polite", choices=["polite", "direct", "concise", "plain"]) |
| parser.add_argument("--weights", type=Path, default=Path("runs/rtx4060/best.pt")) |
| parser.add_argument("--tokenizer", type=Path, default=Path("runs/rtx4060")) |
| parser.add_argument("--device", default="auto", choices=["auto", "cuda", "cpu"]) |
| args = parser.parse_args() |
|
|
| |
| if args.device == "auto": |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| else: |
| device = torch.device(args.device) |
|
|
| run_dir = Path("runs/rtx4060") if (Path("runs/rtx4060") / "tokenizer.json").exists() else Path(".") |
| checkpoint = "best.pt" |
| model, tokenizer = load_model(run_dir, checkpoint, device) |
|
|
| if args.interactive or not args.text: |
| run_interactive(model, tokenizer, device) |
| else: |
| req = [{"source": args.text, "task": args.task, "tone": args.tone, "style": args.style}] |
| out = rewrite(model, tokenizer, req, max_source_length=128, max_new_tokens=128)[0] |
| print(f"\nOriginal : {args.text}") |
| print(f"Settings : task={args.task}, tone={args.tone}, style={args.style}") |
| print(f"Rewritten: {out}\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
