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