arbitration_system_V20/transcript_extractor/rectify_transcript.py

import requests
import json
from typing import Optional
import config.config


class RectifyClient_transcript:
    """
    使用 DeepSeek API 整理 OCR 文本，去除无关标记，保留原文内容。
    """

    def __init__(self, api_key: str = None, model: str = "deepseek-chat",
                 temperature: float = 0.0, max_tokens: int = 4000):
        """
        初始化 OCR 文本清理器。

        Args:
            api_key (str): DeepSeek API 密钥
            model (str): 使用的模型名称，默认为 "deepseek-chat"
            temperature (float): 温度参数，控制随机性，0 表示最确定
            max_tokens (int): 最大输出 token 数，根据文本长度调整
        """
        self.api_key = config.config.DEEPSEEK_API
        self.model = model
        self.temperature = temperature
        self.max_tokens = max_tokens
        self.api_url = "https://api.deepseek.com/v1/chat/completions"

        # 系统提示，明确任务要求
        self.system_prompt = """你是一个文本整理助手。请将用户提供的OCR提取文本整理为纯文本，要求：
1. 去除明显的OCR标记、页码、页眉页脚、行号等无关信息（例如"--- Page X ---"、"Process finished with exit code 0"等）；
2. 保留所有原始文本内容，包括姓名、身份证号、电话号码、统一社会信用代码、金额、日期等数字信息必须与原文完全一致；
3. 不要修改或压缩任何正文文字；
4. 不要使用任何Markdown格式（如#、*、`等），直接输出纯文本。"""

    def clean_text(self, ocr_text: str) -> Optional[str]:
        """
        清理 OCR 文本。

        Args:
            ocr_text (str): OCR 提取的原始文本

        Returns:
            Optional[str]: 清理后的纯文本，如果调用失败则返回 None
        """
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": self.system_prompt},
                {"role": "user", "content": f"请整理以下OCR文本：\n\n{ocr_text}"}
            ],
            "temperature": self.temperature,
            "max_tokens": self.max_tokens
        }

        try:
            response = requests.post(self.api_url, headers=headers, json=payload)
            response.raise_for_status()
            result = response.json()
            cleaned_text = result['choices'][0]['message']['content']
            return cleaned_text
        except requests.exceptions.RequestException as e:
            print(f"HTTP 请求失败: {e}")
        except KeyError as e:
            print(f"API 返回格式异常: {e}")
        except Exception as e:
            print(f"未知错误: {e}")
        return None

    def clean_file(self, input_file: str, output_file: str, encoding: str = "utf-8") -> bool:
        """
        清理文件中的 OCR 文本，并将结果保存到另一个文件。

        Args:
            input_file (str): 输入文件路径（包含 OCR 原始文本）
            output_file (str): 输出文件路径（保存清理后的文本）
            encoding (str): 文件编码，默认为 utf-8

        Returns:
            bool: 是否成功处理
        """
        try:
            with open(input_file, "r", encoding=encoding) as f:
                ocr_text = f.read()

            cleaned = self.clean_text(ocr_text)
            if cleaned is None:
                return False

            with open(output_file, "w", encoding=encoding) as f:
                f.write(cleaned)
            return True
        except IOError as e:
            print(f"文件读写错误: {e}")
            return False


# 使用示例
if __name__ == "__main__":
    # 请替换为你的 DeepSeek API 密钥
    API_KEY = "your-deepseek-api-key"

    # 创建清理器实例
    cleaner = RectifyClient_transcript(api_key=API_KEY)

    # 方式1：直接清理字符串
    ocr_sample = """--- Page 1 ---
苏州市劳动人事争议仲裁委员会
庭审笔录

苏劳人仲案字〔2023〕第1号
2023年1月19日13时30分
...（省略）...
Process finished with exit code 0"""

    cleaned = cleaner.clean_text(ocr_sample)
    if cleaned:
        print("清理后的文本：")
        print(cleaned)

    # 方式2：处理文件
    # success = cleaner.clean_file("ocr_output.txt", "cleaned_output.txt")
    # if success:
    #     print("文件处理成功")