proposa-ai/pdf_table_to_docx/text_scorer.py

"""
text_scorer.py - 基于语言模型的文本完整性评分模块

使用中文 GPT-2 模型（ONNX 格式）计算文本的自然度评分，
用于判断：
1. 两个文本片段是否应该拼接（跨页内容延续检测）
2. 文本中的换行符是否需要去除
"""

import logging
import os
from typing import Optional

import numpy as np
import onnxruntime as ort
from tokenizers import Tokenizer

logger = logging.getLogger(__name__)

# 模型路径
MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)),
                         "models", "gpt2-chinese-cluecorpussmall-onnx")
MODEL_PATH = os.path.join(MODEL_DIR, "model.onnx")
TOKENIZER_PATH = os.path.join(MODEL_DIR, "tokenizer.json")

# 特殊 token ID（BERT WordPiece）
CLS_ID = 101
SEP_ID = 102
PAD_ID = 0

# 最大序列长度
MAX_LEN = 64


class TextScorer:
    """使用中文 GPT-2 计算文本自然度评分的单例封装。"""

    _instance = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance

    def __init__(self):
        if self._initialized:
            return
        self._initialized = True
        self._session = None
        self._tokenizer = None

    def _ensure_loaded(self):
        if self._session is not None:
            return
        if not os.path.exists(MODEL_PATH):
            raise FileNotFoundError(
                f"模型文件不存在: {MODEL_PATH}\n"
                f"请确认 models/gpt2-chinese-cluecorpussmall-onnx/ 已下载"
            )
        self._session = ort.InferenceSession(MODEL_PATH)
        self._tokenizer = Tokenizer.from_file(TOKENIZER_PATH)
        logger.info(f"文本评分模型已加载: {MODEL_PATH}")

    def score(self, text: str) -> float:
        """计算文本的自然度评分（0~1），越高越自然完整。

        基于困惑度（Perplexity）转换：score = 1 / (1 + PPL/10)
        """
        self._ensure_loaded()
        if not text or not text.strip():
            return 1.0

        tokens = self._tokenizer.encode(text.strip())
        ids = tokens.ids

        if len(ids) > MAX_LEN - 2:
            ids = ids[:MAX_LEN - 2]

        input_ids = [CLS_ID] + ids + [SEP_ID]
        seq_len = len(input_ids)
        padding = [PAD_ID] * (MAX_LEN - seq_len)
        input_ids_padded = input_ids + padding
        attn_mask = [1] * seq_len + [0] * (MAX_LEN - seq_len)

        inputs = {
            'input_ids': np.array([input_ids_padded], dtype=np.int64),
            'attention_mask': np.array([attn_mask], dtype=np.int64),
        }
        logits = self._session.run(None, inputs)[0]

        # 计算每个位置的交叉熵损失（预测下一个 token）
        losses = []
        for i in range(seq_len - 1):
            pred = logits[0, i, :]
            true_id = input_ids[i + 1]
            # softmax + cross-entropy
            pred_max = np.max(pred)
            exp_pred = np.exp(pred - pred_max)
            probs = exp_pred / np.sum(exp_pred)
            loss = -np.log(max(probs[true_id], 1e-10))
            losses.append(loss)

        avg_loss = np.mean(losses) if losses else 0
        ppl = np.exp(avg_loss)
        return float(1.0 / (1.0 + ppl / 10.0))

    def is_continuation(self, text1: str, text2: str,
                        threshold: float = 1.5) -> bool:
        """判断 text2 是否是 text1 的内容延续。

        比较 text1、text2 和拼接后文本的评分。
        如果拼接后评分显著高于各自评分，说明是延续关系。

        Args:
            text1: 前半段文本
            text2: 后半段文本
            threshold: 拼接提升倍率阈值（默认 1.5 倍）

        Returns:
            bool: 是否构成内容延续
        """
        if not text1 or not text2:
            return False
        if len(text1.strip()) < 2 or len(text2.strip()) < 2:
            return False

        s1 = self.score(text1)
        s2 = self.score(text2)
        sc = self.score(text1.strip() + text2.strip())

        # 拼接后评分 > 各自最大评分的 threshold 倍
        best_individual = max(s1, s2)
        if best_individual <= 0:
            return False
        improvement = sc / best_individual

        # 同时检查：拼接后不低于单独评分，且明显提升
        return sc > best_individual and improvement >= threshold

    def clean_newlines(self, text: str,
                       threshold: float = 1.1) -> str:
        """去除文本中对自然度有负面影响的换行符。

        比较原文和去除 \\n 后的评分，若去除后评分提升超过 threshold，
        则返回去除换行符的版本。

        Args:
            text: 原始文本
            threshold: 评分提升阈值（默认 1.1 倍）

        Returns:
            str: 清理后的文本
        """
        if not text or '\n' not in text:
            return text

        cleaned = text.replace('\n', '')
        if cleaned == text:
            return text

        s_orig = self.score(text)
        s_clean = self.score(cleaned)

        if s_clean > s_orig * threshold:
            logger.debug(f"去除换行符: '{text[:30]}' -> '{cleaned[:30]}'"
                         f" (评分 {s_orig:.4f} -> {s_clean:.4f})")
            return cleaned

        return text


# 全局单例
_scorer = None


def get_scorer() -> TextScorer:
    """获取全局 TextScorer 单例。"""
    global _scorer
    if _scorer is None:
        _scorer = TextScorer()
    return _scorer