utils/title_scoring.py

from transformers import pipeline
import torch

def score_article_titles(title_list, model_name="mrm8488/bert-mini-finetuned-age_news-classification"):
    """
    为文章标题列表打分，返回对应的分数列表
    
    参数：
        title_list: 文章标题列表（list[str]）
        model_name: Hugging Face 免费模型名称，默认使用轻量的BERT微调模型
    
    返回：
        score_list: 分数列表（list[float]），分数范围0-1，越高代表标题质量越好
    """
    # 检查输入是否为空
    if not title_list:
        return []
    
    # 设置设备（优先使用GPU，没有则用CPU）
    device = 0 if torch.cuda.is_available() else -1
    
    # 初始化评分管道（使用情感/分类模型，适配标题评分场景）
    try:
        # 加载评分管道，返回标签和置信度
        classifier = pipeline(
            "text-classification",
            model=model_name,
            device=device,
            top_k=None  # 返回所有类别的置信度
        )
    except Exception as e:
        print(f"加载模型失败：{e}")
        print("将使用默认的distilbert-base-uncased-emotion模型")
        classifier = pipeline(
            "text-classification",
            model="distilbert-base-uncased-emotion",
            device=device,
            top_k=None
        )
    
    # 对每个标题打分
    score_list = []
    for title in title_list:
        # 过滤空标题
        if not title.strip():
            score_list.append(0.0)
            continue
        
        # 获取模型输出
        pipeline_output = classifier(title)
        
        # 兼容不同版本 transformers 的输出格式
        if isinstance(pipeline_output, list):
            # 如果是嵌套列表 [[{...}, {...}]]，取第一个元素
            if len(pipeline_output) > 0 and isinstance(pipeline_output[0], list):
                results = pipeline_output[0]
            # 否则如果是单层列表 [{...}]，直接作为 results
            else:
                results = pipeline_output
        else:
            # 单个字典的情况
            results = [pipeline_output]

        # 计算综合分数
        total_score = max([item["score"] for item in results])
        # 归一化到0-1区间
        normalized_score = round(max(0.0, min(1.0, total_score)), 4)
        score_list.append(normalized_score)
    
    return score_list