import os import json import tempfile import subprocess from pathlib import Path # 统一把 HOME 指到 /tmp os.environ["HOME"] = "/tmp" Path("/tmp").mkdir(parents=True, exist_ok=True) # 再确保所有 streamlit 相关路径也指向 /tmp os.environ["XDG_CACHE_HOME"] = "/tmp" os.environ["STREAMLIT_CACHE_DIR"] = "/tmp" os.environ["STREAMLIT_GLOBAL_DATA_DIR"] = "/tmp" os.environ["STREAMLIT_RUNTIME_DIR"] = "/tmp" os.environ["STREAMLIT_CONFIG_DIR"] = "/tmp/.streamlit" Path("/tmp/.streamlit").mkdir(parents=True, exist_ok=True) # 可选：关掉使用统计 os.environ["STREAMLIT_BROWSER_GATHER_USAGE_STATS"] = "false" os.environ.setdefault("STREAMLIT_SERVER_ENABLE_CORS", "false") os.environ.setdefault("STREAMLIT_SERVER_ENABLE_XSRF_PROTECTION", "false") import streamlit as st import langextract as lx # pip: langextract[openai] # 你的原始页面设置 st.set_page_config(page_title="LangmyOCR (Streamlit)", layout="wide") st.title("LangmyOCR: OCRmyPDF + LangExtract (Streamlit Demo)") st.caption("先 OCR，后（可选）结构化抽取与交互式复核。数据仅用于会话处理。") # ---------------- Utilities ---------------- def has_bin(name: str) -> bool: return subprocess.run(["bash", "-lc", f"command -v {name} >/dev/null 2>&1"]).returncode == 0 def run_ocr(pdf_file, langs: str, rotate_pages: bool, deskew: bool, clean: bool, optimize_level: int, force_ocr: bool, skip_text: bool, export_sidecar: bool): if pdf_file is None: st.error("请先上传 PDF。") return None, None, None if not has_bin("ocrmypdf"): st.error("系统未检测到 ocrmypdf，可检查 Docker/依赖安装。") return None, None, None # 修复：重置文件指针到开头，然后读取内容 try: pdf_file.seek(0) # 重置文件指针 pdf_content = pdf_file.read() if not pdf_content: st.error("PDF 文件内容为空。") return None, None, None except Exception as e: st.error(f"读取 PDF 文件失败：{e}") return None, None, None # 保存上传文件到临时路径 with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp: tmp.write(pdf_content) # 使用读取到的内容 in_path = Path(tmp.name) work = Path(tempfile.mkdtemp(prefix="ocr_")) out_pdf = work / "output_ocr.pdf" sidecar = work / "out.txt" cmd = ["ocrmypdf", "-l", langs, str(in_path), str(out_pdf)] # 选项（插入到命令靠前位置，保持可读） if rotate_pages: cmd.insert(1, "--rotate-pages") if deskew: cmd.insert(1, "--deskew") if clean: cmd.insert(1, "--clean") cmd[1:1] = ["--optimize", str(optimize_level)] if skip_text: cmd.insert(1, "--skip-text") if force_ocr: cmd.insert(1, "--force-ocr") if export_sidecar: cmd[1:1] = ["--sidecar", str(sidecar)] cmd[1:1] = ["--output-type", "pdfa"] # 归档友好 with st.status("正在执行 OCR …", expanded=False) as s: proc = subprocess.run(cmd, capture_output=True, text=True) if proc.returncode != 0: s.update(label="OCR 失败", state="error") st.error(f"OCR 失败（退出码 {proc.returncode}）") st.code(proc.stderr[-2000:], language="bash") return None, None, None s.update(label="OCR 完成", state="complete") preview = "" sidecar_path = None if export_sidecar and sidecar.exists(): sidecar_path = str(sidecar) try: preview = sidecar.read_text(encoding="utf-8", errors="ignore")[:3000] except Exception: preview = "(sidecar 文本预览读取失败)" return str(out_pdf), sidecar_path, preview def run_extract(sidecar_text: str, provider: str, model_id: str, prompt: str): if not sidecar_text: return None, None, "没有可供抽取的文本。" if provider == "None": return None, None, "未选择模型，跳过抽取。" # 1) 读取 Key，并统一默认打开 fence_output fence_output = True # << 对 Gemini 也打开 use_schema_constraints = False # 先不启 Schema（必要时再开） if provider == "Gemini": api_key = os.environ.get("LANGEXTRACT_API_KEY") if not api_key: return None, None, "未检测到 Gemini API Key（LANGEXTRACT_API_KEY）。" elif provider == "OpenAI": api_key = os.environ.get("OPENAI_API_KEY") if not api_key: return None, None, "未检测到 OpenAI API Key（OPENAI_API_KEY）。" else: return None, None, "未知的 provider。" # 2) 收紧提示语（覆盖面向法律的 schema），严格要求“只返回 JSON 数组” strict_prompt = ( "You are an information extraction engine. " "Extract legal entities, events, relationships, and evidence anchors from the input text. " "Return ONLY a JSON array, no prose, no markdown, no comments. " "Schema per item: {" "\"class\": one of [\"party\",\"event\",\"date\",\"relation\",\"evidence\"], " "\"text\": string (exact span), " "\"attributes\": object (key-value), " "\"source_hint\": string (optional page/line) " "}." ) # 3) 精简可运行的 few-shot（与法律场景贴近） examples = [ lx.data.ExampleData( text="On 15 February 2022, Dr Gavin Soo completed a medicolegal report to Walker Law Group.", extractions=[ lx.data.Extraction( extraction_class="party", extraction_text="Walker Law Group", attributes={"role": "law_firm"}, ), lx.data.Extraction( extraction_class="event", extraction_text="completed a medicolegal report", attributes={"actor": "Dr Gavin Soo"}, ), lx.data.Extraction( extraction_class="date", extraction_text="15 February 2022", attributes={} ), ], ) ] # 4) 先跑一次；若解析失败，再以更强硬提示重试一次 work = Path(tempfile.mkdtemp(prefix="lx_")) jsonl_path = work / "extractions.jsonl" html_path = work / "review.html" raw_path1 = work / "raw_attempt1.txt" raw_path2 = work / "raw_attempt2.txt" def _try_extract(prompt_text): # LangExtract 没有公开 raw 输出参数，我们用 try/except 捕获并让其保存在日志（同时缩短输入验证） return lx.extract( text_or_documents=sidecar_text[:15000], # 先限长，避免超长触发安全策略 prompt_description=prompt_text.strip(), examples=examples, model_id=model_id.strip(), api_key=api_key, fence_output=fence_output, use_schema_constraints=use_schema_constraints, ) with st.status("正在进行结构化抽取 …", expanded=False) as s: try: result = _try_extract(strict_prompt) except Exception as e1: # 第一次失败：很可能是返回了非 JSON。我们把提示再加强，强调 “only JSON array” hard_prompt = strict_prompt + " Output must be a compact JSON array. Do not include any other text." try: result = _try_extract(hard_prompt) except Exception as e2: s.update(label="抽取失败", state="error") return None, None, f"LangExtract 抽取失败：{e2}" # 保存结果并可视化 try: lx.io.save_annotated_documents([result], output_name=str(jsonl_path)) html_content = lx.visualize(str(jsonl_path)) html_path.write_text(html_content, encoding="utf-8") except Exception as e: s.update(label="可视化失败", state="error") return None, None, f"可视化失败：{e}" s.update(label="抽取完成", state="complete") return str(html_path), str(jsonl_path), "抽取成功。" # ---------------- UI ---------------- with st.sidebar: st.header("参数") # 用 form 把"上传 + 参数 + 提交"打包，避免按钮重跑导致 file_uploader 丢值 with st.form("run_form", clear_on_submit=False): pdf = st.file_uploader("上传扫描 PDF", type=["pdf"], accept_multiple_files=False, key="pdf_uploader") langs = st.text_input("OCR 语言（Tesseract 语法）", value="eng+chi_sim") col_a, col_b, col_c = st.columns(3) with col_a: rotate_pages = st.checkbox("自动旋转校正", value=True) with col_b: deskew = st.checkbox("去偏斜", value=True) with col_c: clean = st.checkbox("清理底噪/污渍", value=True) optimize_level = st.select_slider("优化级别", options=[0,1,2], value=1) skip_text = st.checkbox("跳过已有文本层 (--skip-text)", value=True) force_ocr = st.checkbox("强制重做文本层 (--force-ocr) [谨慎]", value=False) export_sidecar = st.checkbox("导出 sidecar 文本", value=True) st.markdown("---") provider = st.selectbox("抽取提供方", ["None", "Gemini", "OpenAI"], index=0) model_id = st.text_input("模型 ID", value="gemini-2.5-flash") prompt = st.text_area( "抽取任务描述（建议按你的法律场景定制）", value=("Extract legal entities, events, relationships, and evidence anchors. " "Return JSON objects with fields: {party, role, event, date, relation, citation, quote}. " "Preserve exact source spans for traceability."), height=160, ) submitted = st.form_submit_button("运行 OCR（+可选抽取）", type="primary") col1, col2 = st.columns([1,1]) with col1: st.subheader("OCR 结果") ocr_pdf_slot = st.empty() sidecar_slot = st.empty() preview_slot = st.empty() with col2: st.subheader("抽取与复核") html_slot = st.empty() jsonl_slot = st.empty() status_slot = st.empty() # 辅助：显示文件已被接收（提交前就可见，便于确认） if "pdf_uploader" in st.session_state and st.session_state["pdf_uploader"]: st.sidebar.success(f"已选择：{st.session_state['pdf_uploader'].name} " f"({st.session_state['pdf_uploader'].size/1024:.1f} KB)") if submitted: # 添加调试信息 if pdf is None: st.error("PDF 为 None - 检查文件上传") else: st.info(f"PDF 文件信息：名称={pdf.name}, 大小={pdf.size} bytes") out_pdf, sidecar_path, preview = run_ocr( pdf, langs, rotate_pages, deskew, clean, optimize_level, force_ocr, skip_text, export_sidecar ) if out_pdf: with open(out_pdf, "rb") as f: ocr_pdf_slot.download_button("下载 OCR 后 PDF", f, file_name="output_ocr.pdf") if sidecar_path: with open(sidecar_path, "rb") as f: sidecar_slot.download_button("下载 sidecar 文本", f, file_name="out.txt") preview_slot.text_area("sidecar 文本预览（前 3000 字）", value=preview, height=240) if sidecar_path and provider != "None": txt = Path(sidecar_path).read_text(encoding="utf-8", errors="ignore") html_path, jsonl_path, status = run_extract(txt, provider, model_id, prompt) status_slot.info(status) if html_path and Path(html_path).exists(): html_content = Path(html_path).read_text(encoding="utf-8", errors="ignore") st.components.v1.html(html_content, height=650, scrolling=True) if jsonl_path and Path(jsonl_path).exists(): with open(jsonl_path, "rb") as f: jsonl_slot.download_button("下载抽取结果 JSONL", f, file_name="extractions.jsonl")