File size: 3,719 Bytes
2651102 | 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 | """
Diffutslator 主入口
基于扩散模型的中英互译系统
"""
import os
import sys
import argparse
def main():
parser = argparse.ArgumentParser(
description="Diffutslator - 基于扩散模型的翻译系统",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
示例:
# 快速验证训练
python main.py train --quick
# 完整训练
python main.py train --full
# 从检查点恢复训练
python main.py train --resume checkpoints/epoch_5.pt
# 交互式翻译
python main.py translate
# 翻译单个句子
python main.py translate --text "你好世界" --zh
# 使用更多DDIM步数
python main.py translate --text "Hello world" --en --ddim-steps 100
"""
)
subparsers = parser.add_subparsers(dest="command", help="命令")
# 训练命令
train_parser = subparsers.add_parser("train", help="训练模型")
train_parser.add_argument("--quick", action="store_true", help="快速验证模式")
train_parser.add_argument("--full", action="store_true", help="完整训练模式")
train_parser.add_argument("--samples", type=int, default=None, help="使用的数据量")
train_parser.add_argument("--epochs", type=int, default=None, help="训练轮数")
train_parser.add_argument("--batch-size", type=int, default=None, help="批量大小")
train_parser.add_argument("--resume", type=str, default=None, help="恢复训练的检查点")
# 翻译命令
translate_parser = subparsers.add_parser("translate", help="翻译文本")
translate_parser.add_argument("--checkpoint", type=str, default=None, help="检查点路径")
translate_parser.add_argument("--text", type=str, default=None, help="要翻译的文本")
translate_parser.add_argument("--zh", action="store_true", help="输入是中文")
translate_parser.add_argument("--en", action="store_true", help="输入是英文")
translate_parser.add_argument("--interactive", "-i", action="store_true", help="交互模式")
translate_parser.add_argument("--quiet", "-q", action="store_true", help="安静模式")
translate_parser.add_argument("--ddim-steps", type=int, default=50, help="DDIM步数")
args = parser.parse_args()
if args.command == "train":
# 导入并运行训练
from train import main as train_main
sys.argv = ["train.py"]
if args.quick:
sys.argv.append("--quick")
if args.full:
sys.argv.append("--full")
if args.samples:
sys.argv.extend(["--samples", str(args.samples)])
if args.epochs:
sys.argv.extend(["--epochs", str(args.epochs)])
if args.batch_size:
sys.argv.extend(["--batch-size", str(args.batch_size)])
if args.resume:
sys.argv.extend(["--resume", args.resume])
train_main()
elif args.command == "translate":
# 导入并运行推理
from inference import main as inference_main
sys.argv = ["inference.py"]
if args.checkpoint:
sys.argv.extend(["--checkpoint", args.checkpoint])
if args.text:
sys.argv.extend(["--text", args.text])
if args.zh:
sys.argv.append("--zh")
if args.en:
sys.argv.append("--en")
if args.interactive:
sys.argv.append("--interactive")
if args.quiet:
sys.argv.append("--quiet")
if args.ddim_steps:
sys.argv.extend(["--ddim-steps", str(args.ddim_steps)])
inference_main()
else:
parser.print_help()
if __name__ == "__main__":
main()
|