Update umt5_summary.py

d6620b5e · wanglch · 3d0ce953 · d6620b5e
Commit d6620b5e authored May 15, 2024 by wanglch
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 0 deletions

umt5_summary.py umt5_summary.py +1 -0

No files found.
--- a/umt5_summary.py
+++ b/umt5_summary.py
@@ -11,6 +11,7 @@ print(f'Using {device} device')

 # 根据用户路径修改对应路径即可
 model_checkpoint = "/umt5/umt5_base"
+# 训练权重需要预训练后载入
 trained_model_weights = '/umt5/saves/train_dtk_weights/epoch_1_valid_rouge_23.4347_model_dtk_weights.bin'
 tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
 model = AutoModelForSeq2SeqLM.from_pretrained(model_checkpoint)