File size: 1,488 Bytes
a1da864 7c7241f 951bc06 7c7241f a1da864 7c7241f a1da864 951bc06 a1da864 951bc06 a1da864 951bc06 a1da864 951bc06 a1da864 951bc06 a1da864 951bc06 7c7241f a1da864 951bc06 7c7241f 951bc06 a1da864 951bc06 a1da864 951bc06 a1da864 951bc06 | 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 | from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
def summarize_text(text, max_length=300, min_length=50):
"""
使用 Qwen2.5-0.5B-Instruct 生成摘要
:param text: 输入的文章内容（支持中英文）
:param max_length: 摘要最大长度（汉字数）
:param min_length: 摘要最小长度（汉字数）
:return: 生成的摘要
"""
model_name = "Qwen/Qwen2.5-0.5B-Instruct"
# 加载模型和分词器
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto")
model.to(device)
tokenizer = AutoTokenizer.from_pretrained(model_name)
# 构建提示
prompt = f"请将以下文章总结为{max_length}字以内的摘要：\n\n{text}"
messages = [{"role": "user", "content": prompt}]
# 应用聊天模板
text = tokenizer.apply_chat_template(
messages, tokenize=False, add_generation_prompt=True
)
# 生成摘要
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(
**model_inputs, max_new_tokens=max_length + 50, temperature=0.3
)
# 处理生成结果
generated_ids = [
output_ids[len(input_ids) :]
for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
return response
|