# 1. 创建虚拟环境
# conda create -n qwen_sft python=3.12 -y
# conda activate qwen_sft

# 2. 安装核心依赖
# pip install torch torchvision torchaudio
# pip install transformers accelerate sentencepiece modelscope

from modelscope import snapshot_download
from transformers import AutoModelForCausalLM, AutoTokenizer

# 下载基座模型
snapshot_download('Qwen/Qwen2.5-1.5B', cache_dir='./qwen2.5-1.5b')

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    "./qwen2.5-1.5b/Qwen/Qwen2.5-1.5B",
    device_map="auto",
    torch_dtype="auto"
)
tokenizer = AutoTokenizer.from_pretrained("./qwen2.5-1.5b/Qwen/Qwen2.5-1.5B")

# 测试样本
ARTICLE = """微软推出新型AI芯片Maia 100，采用台积电5nm工艺，配备120GB HBM3内存，FP8算力达1250 TFLOPS。该芯片将用于Azure云服务，推理性能较前代提升40%，预计2025年Q1量产"""

PROMPT_TEMPLATE = f"""生成科技新闻摘要：
### 原文:
{ARTICLE}
### 摘要:
"""

# 生成摘要
messages = [{"role": "user", "content": PROMPT_TEMPLATE}]
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
output = model.generate(inputs, max_new_tokens=200)
result = tokenizer.decode(output[0], skip_special_tokens=True)
print(result)