wordformat 1.0.0__py3-none-any.whl

wordformat/__init__.py

@@ -0,0 +1,11 @@
+#! /usr/bin/env python
+# @Time    : 2025/12/22 21:47
+# @Author  : afish
+# @File    : __init__.py.py
+from wordformat.set_style import auto_format_thesis_document
+from wordformat.set_tag import set_tag_main
+
+__all__ = [
+    "auto_format_thesis_document",
+    "set_tag_main"
+]

wordformat/agent/__init__.py

@@ -0,0 +1,4 @@
+#! /usr/bin/env python
+# @Time    : 2026/1/11 17:38
+# @Author  : afish
+# @File    : __init__.py.py

wordformat/agent/message.py

@@ -0,0 +1,79 @@
+#! /usr/bin/env python
+# @Time    : 2026/1/11 17:39
+# @Author  : afish
+# @File    : message.py
+import threading
+
+
+class MessageManager:
+    def __init__(self):
+        self._messages = []
+        self._instance_lock = threading.Lock()  # 实例级操作锁
+
+    def add_message(self, role: str, content: str):
+        """添加角色信息"""
+        with self._instance_lock:
+            self._messages.append({"role": role, "content": content})
+
+    def add_dict_message(self, content):
+        """保存对话历史或上下文信息"""
+        with self._instance_lock:
+            self._messages.append(content)
+
+    def add_user_message(self, content: str):
+        """添加角色信息"""
+        self.add_dict_message({"role": "user", "content": content})
+
+    def add_assistant_message(self, msg: str | dict):
+        """安全添加AI消息"""
+        if isinstance(msg, dict):
+            self.add_dict_message(msg)
+            return
+        content = msg.content if hasattr(msg, "content") else msg
+        tool_calls = getattr(msg, "tool_calls", None)
+        message = {"role": "assistant", "content": content}
+        if tool_calls:
+            message["tool_calls"] = [
+                {
+                    "id": call.id,
+                    "type": call.type,
+                    "function": {
+                        "name": call.function.name,
+                        "arguments": call.function.arguments,
+                    },
+                }
+                for call in tool_calls
+            ]
+
+        self.add_dict_message(message)
+
+    def add_tool_message(self, content: str, tool_call_id):
+        """添加工具信息"""
+        self.add_dict_message(
+            {"role": "tool", "content": content, "tool_call_id": tool_call_id}
+        )
+
+    def add_system_message(self, content: str):
+        """添加系统信息"""
+        self.add_dict_message({"role": "system", "content": content})
+
+    def get_messages(self):
+        """获取信息"""
+        with self._instance_lock:
+            return self._messages.copy()
+
+    def reset_messages(self):
+        """清空信息"""
+        with self._instance_lock:
+            self._messages.clear()
+
+    @property
+    def messages(self):
+        """返回信息列表"""
+        with self._instance_lock:
+            return self._messages.copy()
+
+    def clear(self):
+        self._messages = list(
+            filter(lambda msg: msg.get("role") == "system", self._messages)
+        )

wordformat/agent/onnx_infer.py

@@ -0,0 +1,261 @@
+import json
+import os
+import time
+from typing import Dict, List, Optional
+
+import numpy as np
+from loguru import logger
+
+# ===== 全局变量（初始为 None）=====
+_tokenizer: Optional["Tokenizer"] = None  # noqa F821
+_ort_sess: Optional["ort.InferenceSession"] = None  # noqa F821
+_id2label: Optional[Dict[int, str]] = None
+MAX_LENGTH = 128
+
+
+# ===== 路径配置（仅定义，不加载）=====
+def _get_model_paths():
+    from importlib.resources import files
+
+    return {
+        "onnx": str(
+            files("wordformat.data.model").joinpath("bert_paragraph_classifier.onnx")
+        ),
+        "tokenizer": str(files("wordformat.data.model").joinpath("tokenizer.json")),
+        "id2label": str(files("wordformat.data.model").joinpath("id2label.json")),
+    }
+
+
+# ===== 硬件自动选择函数 =====
+def _get_best_onnx_providers() -> List[str]:
+    """自动选择最优推理硬件：CUDA > DirectML > CPU"""
+    try:
+        import onnxruntime as ort
+
+        available = ort.get_available_providers()
+        # 优先级1：NVIDIA显卡（CUDA）
+        if "CUDAExecutionProvider" in available:
+            logger.info("检测到CUDA，优先使用GPU推理")
+            return ["CUDAExecutionProvider"]
+        # 优先级2：Windows核显（Intel/AMD，DirectML）
+        elif "DmlExecutionProvider" in available:
+            logger.info("检测到核显，使用DirectML推理")
+            return ["DmlExecutionProvider"]
+        # 优先级3：纯CPU（兜底）
+        else:
+            logger.info("仅检测到CPU，使用CPU多核推理")
+            return ["CPUExecutionProvider"]
+    except Exception as e:
+        logger.warning(f"硬件检测失败，降级为CPU：{e}")
+        return ["CPUExecutionProvider"]
+
+
+# ===== 模型加载函数 =====
+def _load_model():
+    global _tokenizer, _ort_sess, _id2label
+    if _tokenizer is not None:
+        return
+
+    import onnxruntime as ort
+    from tokenizers import Tokenizer
+
+    paths = _get_model_paths()
+    logger.info(f"首次调用，正在加载模型：{paths['onnx']}")
+
+    # 1. 加载Tokenizer（原有逻辑保留）
+    _tokenizer = Tokenizer.from_file(paths["tokenizer"])
+
+    # 2. 优化ONNX推理器配置
+    ort_options = ort.SessionOptions()
+    ort_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    cpu_core_num = os.cpu_count() or 4
+    ort_options.intra_op_num_threads = cpu_core_num
+    ort_options.inter_op_num_threads = cpu_core_num
+    ort_options.log_severity_level = 3
+    ort_options.enable_cpu_mem_arena = True
+    ort_options.enable_mem_pattern = True
+    ort_options.enable_mem_reuse = True
+    ort_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+
+    # 3. 加载ONNX模型（自动适配硬件）
+    providers = _get_best_onnx_providers()
+
+    try:
+        _ort_sess = ort.InferenceSession(
+            paths["onnx"], sess_options=ort_options, providers=providers
+        )
+
+    except Exception as e:
+        logger.warning(f"最优硬件加载失败，降级为CPU：{e}")
+        _ort_sess = ort.InferenceSession(
+            paths["onnx"], sess_options=ort_options, providers=["CPUExecutionProvider"]
+        )
+    # 4. 加载id2label（原有逻辑保留）
+    with open(paths["id2label"], encoding="utf-8") as f:
+        _id2label = {int(k): v for k, v in json.load(f).items()}
+
+
+def onnx_single_infer(text: str) -> dict:
+    if _tokenizer is None:
+        _load_model()
+
+    # 确保类型安全（静态检查友好）
+    assert _tokenizer is not None
+    assert _ort_sess is not None
+    assert _id2label is not None
+
+    # 使用官方 tokenizer 编码
+    encoded = _tokenizer.encode(text, add_special_tokens=True)
+    input_ids = encoded.ids
+    attention_mask = encoded.attention_mask
+    token_type_ids = encoded.type_ids
+
+    # 截断 & 补全
+    if len(input_ids) > MAX_LENGTH:
+        input_ids = input_ids[:MAX_LENGTH]
+        attention_mask = attention_mask[:MAX_LENGTH]
+        token_type_ids = token_type_ids[:MAX_LENGTH]
+    else:
+        pad_len = MAX_LENGTH - len(input_ids)
+        input_ids += [0] * pad_len
+        attention_mask += [0] * pad_len
+        token_type_ids += [0] * pad_len
+
+    onnx_input = {
+        "input_ids": np.array([input_ids], dtype=np.int64),
+        "attention_mask": np.array([attention_mask], dtype=np.int64),
+        "token_type_ids": np.array([token_type_ids], dtype=np.int64),
+    }
+    # 超时控制
+    try:
+        start = time.time()
+        logits = _ort_sess.run(["logits"], onnx_input)[0]
+        logger.debug(f"单条推理耗时：{time.time() - start:.4f}s")
+    except Exception as e:
+        logger.error(f"单条推理失败：{e}")
+        return {"label": "", "score": 0.0}
+
+    logits = logits[0]
+    logits -= np.max(logits)
+    probs = np.exp(logits) / np.sum(np.exp(logits))
+    pred_id = int(np.argmax(probs))
+    pred_prob = round(float(probs[pred_id]), 4)
+
+    return {"label": _id2label[pred_id], "score": pred_prob}
+
+
+def onnx_batch_infer(texts: list[str]) -> list[dict]:
+    """
+    ONNX模型批量文本推理
+    :param texts: 待分类的论文段落文本列表
+    :return: 每条文本的预测结果列表，每个元素为字典（同单条推理格式）
+    """
+    global _tokenizer, _ort_sess, _id2label
+
+    if not texts:
+        return []
+
+    if _tokenizer is None:
+        _load_model()
+
+    # 确保类型安全（静态检查友好）
+    assert _tokenizer is not None
+    assert _ort_sess is not None
+    assert _id2label is not None
+
+    # ===== 批量预处理（向量化替代循环拼接）=====
+    batch_size = len(texts)
+    # 初始化批量数组（直接预分配内存，避免循环append）
+    batch_input_ids = np.zeros((batch_size, MAX_LENGTH), dtype=np.int64)
+    batch_attention_mask = np.zeros((batch_size, MAX_LENGTH), dtype=np.int64)
+    batch_token_type_ids = np.zeros((batch_size, MAX_LENGTH), dtype=np.int64)
+
+    # 批量编码（保留逻辑，但优化赋值方式）
+    for idx, text in enumerate(texts):
+        encoded = _tokenizer.encode(text, add_special_tokens=True)
+        input_ids = encoded.ids
+        attention_mask = encoded.attention_mask
+        token_type_ids = encoded.type_ids
+
+        # 截断
+        seq_len = min(len(input_ids), MAX_LENGTH)
+        # 直接赋值（比append+拼接快）
+        batch_input_ids[idx, :seq_len] = input_ids[:seq_len]
+        batch_attention_mask[idx, :seq_len] = attention_mask[:seq_len]
+        batch_token_type_ids[idx, :seq_len] = token_type_ids[:seq_len]
+
+    # ===== 批量推理（超时控制+性能监控）=====
+    onnx_input = {
+        "input_ids": batch_input_ids,
+        "attention_mask": batch_attention_mask,
+        "token_type_ids": batch_token_type_ids,
+    }
+
+    try:
+        start = time.time()
+        # 批量推理（一次性处理，减少ONNX调用开销）
+        logits = _ort_sess.run(["logits"], onnx_input)[0]  # shape: [batch, num_classes]
+        infer_time = time.time() - start
+        logger.info(
+            f"批量推理完成 | 批次大小：{batch_size} | 耗时：{infer_time:.4f}s | 单条耗时：{infer_time / batch_size:.4f}s"  # noqa E501
+        )
+    except Exception as e:
+        logger.error(f"批量推理失败：{e}")
+        # 兜底：返回空结果，避免整体崩溃
+        return [
+            {"text": text, "label": "", "pred_id": -1, "score": 0.0}
+            for text in texts
+        ]
+
+    # ===== 向量化计算概率（替代循环）=====
+    # 数值稳定化（避免exp溢出）
+    logits_stable = logits - np.max(logits, axis=-1, keepdims=True)
+    probs = np.exp(logits_stable) / np.sum(
+        np.exp(logits_stable), axis=-1, keepdims=True
+    )
+    # 批量获取预测ID和概率（向量化操作，比循环快）
+    pred_ids = np.argmax(probs, axis=-1).astype(int)
+    pred_probs = np.round(np.max(probs, axis=-1).astype(float), 4)
+
+    # ===== 结果组装 =====
+    results = []
+    for idx, text in enumerate(texts):
+        pred_id = pred_ids[idx]
+        results.append(
+            {
+                "text": text,
+                "label": _id2label.get(pred_id, ""),
+                "pred_id": pred_id,
+                "score": float(pred_probs[idx]),
+            }
+        )
+
+    return results
+
+
+def safe_batch_infer(texts: list[str], max_batch_size: int = 128) -> list[dict]:
+    """
+    安全批量推理（自动分片，避免超大批次OOM）
+    :param texts: 文本列表
+    :param max_batch_size: 单批次最大数量（默认128，可根据硬件调整）
+    :return: 完整结果列表
+    """
+    if not texts:
+        return []
+
+    start = time.time()
+    results = []
+    total = len(texts)
+    # 分片处理
+    for i in range(0, total, max_batch_size):
+        batch = texts[i : i + max_batch_size]
+        batch_results = onnx_batch_infer(batch)
+        results.extend(batch_results)
+        # 进度日志
+        logger.info(f"已处理 {min(i + max_batch_size, total)}/{total} 条文本")
+
+    total_time = time.time() - start
+    logger.info(
+        f"安全批量推理完成 | 总条数：{total} | 总耗时：{total_time:.4f}s | 平均单条：{total_time / total:.4f}s"  # noqa E501
+    )
+    return results

wordformat/api/__init__.py

@@ -0,0 +1,269 @@
+#! /usr/bin/env python
+# @Time    : 2026/2/5 21:41
+# @Author  : afish
+# @File    : __init__.py
+import os
+from pathlib import Path
+from typing import Optional
+from urllib.parse import quote
+
+from fastapi import Body, FastAPI, File, HTTPException, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from loguru import logger
+from pydantic import BaseModel
+from starlette.responses import FileResponse
+
+# 复用原有项目的核心函数和校验工具
+from wordformat.set_style import auto_format_thesis_document
+from wordformat.set_tag import set_tag_main
+from wordformat.settings import BASE_DIR, SERVER_HOST
+
+# ---------------------- 初始化FastAPI应用 ----------------------
+app = FastAPI(
+    title="学位论文格式自动校验工具-WebAPI",
+    description="基于FastAPI实现，支持generate-json/check-format/apply-format三种模式，兼容原有YAML配置",
+    version="1.0.0",
+    docs_url="/docs",  # Swagger UI接口文档地址（推荐）
+    redoc_url="/redoc",  # ReDoc接口文档地址（备选）
+)
+
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],  # 允许上述域名的跨域请求
+    allow_credentials=True,  # 允许携带cookie（可选，建议开）
+    allow_methods=["*"],  # 允许所有请求方法（GET/POST/PUT等）
+    allow_headers=["*"],  # 允许所有请求头（包括文件上传的头）
+)
+
+# ---------------------- 全局配置 ----------------------
+# 临时文件目录（存储上传的docx/配置文件、生成的json），自动创建
+TEMP_DIR = BASE_DIR / "temp"
+TEMP_DIR.mkdir(parents=True, exist_ok=True)
+logger.info(f"临时文件目录：{TEMP_DIR}")
+# 输出文件目录（存储校验/格式化后的docx）
+OUTPUT_DIR = BASE_DIR / "output"
+OUTPUT_DIR.mkdir(parents=True, exist_ok=True)
+logger.info(f"输出文件目录：{OUTPUT_DIR}")  # 修复原日志笔误
+
+
+# ---------------------- 数据模型（接口参数校验） ----------------------
+class OperationResult(BaseModel):
+    """统一接口返回格式"""
+
+    code: int  # 200=成功，400=参数错误，500=执行失败
+    msg: str  # 结果描述
+    data: Optional[dict] = None  # 附加数据（如文件路径、日志信息）
+
+
+# ---------------------- 核心工具函数 ----------------------
+def save_upload_file(upload_file: UploadFile, save_dir: Path) -> str:
+    """
+    保存上传的文件到指定目录，仅返回【文件绝对路径】（重名自动加_1/_2后缀，避免覆盖）
+    :param upload_file: FastAPI上传的文件对象
+    :param save_dir: 保存目录
+    :return: 文件绝对路径
+    """
+    try:
+        original_filename = upload_file.filename
+        file_suffix = os.path.splitext(original_filename)[-1]
+        original_name_no_suffix = os.path.splitext(original_filename)[0]
+
+        # 重名处理：自动追加数字后缀，防止文件覆盖
+        save_path = os.path.join(save_dir, original_filename)
+        counter = 1
+        while os.path.exists(save_path):
+            save_path = os.path.join(
+                save_dir, f"{original_name_no_suffix}_{counter}{file_suffix}"
+            )
+            counter += 1
+
+        # 保存文件流
+        with open(save_path, "wb") as f:
+            f.write(upload_file.file.read())
+        abs_path = os.path.abspath(save_path)
+        logger.info(f"上传文件已保存：{abs_path}")
+        return abs_path
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"文件保存失败：{str(e)}") from e
+
+
+# ---------------------- 核心API接口 ----------------------
+@app.post(
+    "/generate-json", response_model=OperationResult, summary="生成文档结构JSON文件"
+)
+async def api_generate_json(
+    docx_file: UploadFile = File(..., description="待处理的Word文档（.docx格式）"),  # noqa B008
+    config_file: UploadFile = File(..., description="格式配置YAML文件（必填）"),  # noqa B008
+):
+    """
+    对应原命令行generate-json模式：仅生成JSON，不执行校验/格式化
+    - 上传docx和yaml配置文件，服务端自动生成JSON并返回数据
+    """
+
+    try:
+        filename = docx_file.filename.lower()
+        if not filename.endswith(".docx"):
+            return OperationResult(
+                code=400,
+                msg=f"上传失败：仅支持 .docx式（你当前上传的是：{docx_file.filename}），请转换为docx后重试"
+            )
+
+        # 保存上传文件（仅返回路径，无需提取原名称）
+        docx_path = save_upload_file(docx_file, TEMP_DIR)
+        config_path = save_upload_file(config_file, TEMP_DIR)
+
+        # 执行核心逻辑生成JSON
+        json_data = set_tag_main(
+            docx_path=docx_path, configpath=config_path
+        )
+
+        return OperationResult(
+            code=200,
+            msg="JSON文件生成成功",
+            data={
+                "json_data": json_data,
+                "tips": "可使用该JSON数据调用check-format/apply-format接口",
+            },
+        )
+    except Exception as e:
+        logger.error(f"生成JSON失败：{str(e)}")
+        raise HTTPException(status_code=500, detail=f"生成JSON失败：{str(e)}") from e
+
+
+@app.post(
+    "/check-format",
+    response_model=OperationResult,
+    summary="仅执行格式校验（需先生成JSON）",
+)
+async def api_check_format(
+    docx_file: UploadFile = File(..., description="待校验的Word文档（.docx格式）"),  # noqa B008
+    config_file: UploadFile = File(..., description="格式配置YAML文件（必填）"),  # noqa B008
+    json_data: str = Body(..., description="从/generate-json获取的文档结构JSON数据"),
+):
+    """
+    对应原命令行check-format模式：仅执行格式校验，生成【原文件名+--标注版.docx】
+    - 基于函数返回的真实路径拼接下载链接，保证路径100%匹配
+    """
+    try:
+        # 1. 保存上传文件（仅返回实际路径）
+        docx_path = save_upload_file(docx_file, TEMP_DIR)
+        config_path = save_upload_file(config_file, TEMP_DIR)
+
+        # 2. 执行校验逻辑，获取【函数返回的实际保存文件路径】（核心！）
+        actual_save_path = auto_format_thesis_document(
+            jsonpath=json_data,
+            docxpath=docx_path,
+            configpath=config_path,
+            savepath=OUTPUT_DIR,
+            check=True,  # 仅校验模式
+        )
+
+        # 3. 从真实路径中提取最终文件名（如：1 (1)_2--标注版.docx）
+        final_filename = os.path.basename(actual_save_path)
+        # 4. 拼接正确的下载链接（仅编码文件名，无多余路径）
+        encoded_filename = quote(final_filename)
+        download_url = f"{SERVER_HOST}/download/{encoded_filename}"
+
+        # 5. 记录真实的保存路径日志（与函数内部日志一致）
+        logger.info(f"格式校验完成，结果保存：{actual_save_path}")
+
+        # 6. 返回结果（含真实文件名和下载链接）
+        return OperationResult(
+            code=200,
+            msg="格式校验执行成功",
+            data={
+                "original_docx": docx_file.filename,  # 用户上传的原文件名
+                "final_filename": final_filename,  # 实际保存的最终文件名
+                "download_url": download_url,  # 可直接点击的下载链接
+                "tips": "文件已保存至服务端output目录，点击链接直接下载",
+            },
+        )
+    except Exception as e:
+        logger.error(f"格式校验失败：{str(e)}")
+        raise HTTPException(status_code=500, detail=f"格式校验失败：{str(e)}") from e
+
+
+@app.post(
+    "/apply-format", response_model=OperationResult, summary="执行格式应用/自动格式化"
+)
+async def api_apply_format(
+    docx_file: UploadFile = File(..., description="待格式化的Word文档（.docx格式）"),  # noqa B008
+    config_file: UploadFile = File(..., description="格式配置YAML文件（必填）"),  # noqa B008
+    json_data: str = Body(..., description="从/generate-json获取的文档结构JSON数据"),
+):
+    """
+    对应原命令行apply-format模式：自动应用格式，生成【原文件名+--修改版.docx】
+    - 基于函数返回的真实路径拼接下载链接，彻底解决404问题
+    """
+    try:
+        # 1. 保存上传文件（仅返回实际路径）
+        docx_path = save_upload_file(docx_file, TEMP_DIR)
+        config_path = save_upload_file(config_file, TEMP_DIR)
+
+        # 2. 执行格式化逻辑，获取【函数返回的实际保存文件路径】（核心！）
+        actual_save_path = auto_format_thesis_document(
+            jsonpath=json_data,
+            docxpath=docx_path,
+            configpath=config_path,
+            savepath=OUTPUT_DIR,
+            check=False,  # 格式化模式
+        )
+
+        # 3. 从真实路径提取最终文件名（如：1 (1)_2--修改版.docx）
+        final_filename = os.path.basename(actual_save_path)
+        # 4. 拼接正确下载链接（仅编码文件名，避免路径错误）
+        encoded_filename = quote(final_filename)
+        download_url = f"/download/{encoded_filename}"
+
+        # 5. 记录真实保存路径日志（与函数内部日志完全一致）
+        logger.info(f"文档格式化完成，结果保存：{actual_save_path}")
+
+        # 6. 返回结果（含原文件名、最终文件名、下载链接）
+        return OperationResult(
+            code=200,
+            msg="文档格式化执行成功",
+            data={
+                "original_docx": docx_file.filename,  # 用户上传的原文件名
+                "final_filename": final_filename,  # 服务端实际保存的文件名
+                "download_url": download_url,  # 前端可直接使用的下载链接
+                "tips": "文件已保存至服务端output目录，点击链接直接下载",
+            },
+        )
+    except Exception as e:
+        logger.error(f"文档格式化失败：{str(e)}")
+        raise HTTPException(status_code=500, detail=f"文档格式化失败：{str(e)}") from e
+
+
+@app.get("/download/{filename}", summary="下载格式化/校验后的Word文档")
+def download_file(filename: str):
+    """
+    下载接口：基于真实文件名匹配，增加多层校验，保证下载稳定
+    """
+    try:
+        # 拼接服务端实际文件路径
+        file_path = os.path.join(OUTPUT_DIR, filename)
+        # 校验1：文件是否存在
+        if not os.path.exists(file_path):
+            raise HTTPException(status_code=404, detail="文件不存在或已被删除")
+        # 校验2：是否为有效文件（非文件夹/链接）
+        if not os.path.isfile(file_path):
+            raise HTTPException(status_code=400, detail="请求路径不是有效文件")
+        # 核心修复：增加强制下载的响应头
+        headers = {
+            "Content-Disposition": f"attachment; filename={quote(filename)}",  # 强制下载+编码文件名
+            "Cache-Control": "no-cache",  # 避免缓存问题
+            "Pragma": "no-cache",
+        }
+        # 以附件形式返回，浏览器自动触发下载，指定docx专属MIME类型
+        return FileResponse(
+            file_path,
+            filename=filename,  # 强制指定下载显示的文件名（与实际保存一致）
+            media_type="application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+            headers=headers,
+        )
+    except HTTPException as e:
+        logger.error(str(e))
+    except Exception as e:
+        logger.error(f"文件下载失败：{filename} -> {str(e)}")
+        raise HTTPException(status_code=500, detail=f"文件下载失败：{str(e)}") from e