# 注意：请把该文件放到文件根目录运行
# 下载或更新本地的模型
import os

def download_mT5_model_only():
    """只下载 mT5 摘要模型"""
    model_name = "csebuetnlp/mT5_multilingual_XLSum"
    save_path = "./models/bin/mT5_multilingual_XLSum"
    
    print(f"下载摘要模型: {model_name}")
    print(f"保存到: {save_path}")
    
    # 确保目录存在
    os.makedirs(save_path, exist_ok=True)
    
    # 下载模型
    from transformers import MT5ForConditionalGeneration, MT5Tokenizer
    
    print("下载 mT5 模型中...")
    model = MT5ForConditionalGeneration.from_pretrained(model_name)
    print("下载 mT5 分词器中...")
    tokenizer = MT5Tokenizer.from_pretrained(model_name)
    
    print("保存到本地...")
    model.save_pretrained(save_path)
    tokenizer.save_pretrained(save_path)
    
    print(f"✅ mT5 模型已保存到: {save_path}")
    return model, tokenizer

# 在 __main__ 中可以选择调用哪个
if __name__ == "__main__":
    # 只下载摘要模型
    download_mT5_model_only()
    
    # 或者下载所有模型
    # download_all_models()