File size: 4,805 Bytes
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 | """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
# Ensure src is on path
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()
# Determine device
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()
|