import os import re os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE" import gradio as gr import time import numpy as np import io from PIL import Image from dotenv import load_dotenv load_dotenv() import easyocr import faiss import random import requests from openai import OpenAI # 全局会话存储 session_data = {} # session_id -> {"chunks": [...], "index": faiss_index, "history": [...]} reader = easyocr.Reader(['ch_sim', 'en'], gpu=False) def embed_texts(texts): # 这里可替换为真实嵌入接口调用 return [np.random.rand(768).tolist() for _ in texts] def safe_call_model(prompt, session_id, model_list, history): for model in model_list: try: if model == "openai": client = OpenAI(api_key=os.getenv("OPENAI_API_KEY")) response = client.chat.completions.create( model="gpt-3.5-turbo", messages=[ {"role": "system", "content": "你是一个多模态文档理解助手。"}, {"role": "user", "content": prompt} ], temperature=0.5 ) answer = response.choices[0].message.content return { "answer": answer, "model": "gpt-3.5-turbo", "elapsed": 1.2, "cost": 0.002 } elif model == "claude": print("当前使用的ANTHROPIC_API_KEY:", os.getenv("ANTHROPIC_API_KEY")) headers = { "x-api-key": os.getenv("ANTHROPIC_API_KEY"), "content-type": "application/json", "anthropic-version": "2023-06-01" } payload = { "model": "claude-3-sonnet-20240229", "max_tokens": 1024, "temperature": 0.5, "messages": [{"role": "user", "content": prompt}] } response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=payload) if response.status_code == 429: print("⚠️ Claude API 限速，2秒后重试一次...") time.sleep(2) response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=payload) if response.status_code != 200: raise Exception(f"Claude API 请求失败：{response.status_code} {response.text}") data = response.json() if "content" in data and isinstance(data["content"], list): answer = data["content"][0].get("text", "⚠️ Claude返回内容结构异常") else: raise Exception(f"Claude响应格式无效：{data}") return { "answer": answer, "model": "claude-3-sonnet-20240229", "elapsed": 1.2, "cost": 0.002 } elif model == "deepseek": headers = { "Authorization": f"Bearer {os.getenv('DEEPSEEK_API_KEY')}", "Content-Type": "application/json" } payload = { "model": "deepseek-chat", "messages": [ {"role": "system", "content": "你是一个多模态文档理解助手。"}, {"role": "user", "content": prompt} ], "temperature": 0.5 } response = requests.post( os.getenv("DEEPSEEK_BASE_URL"), headers=headers, json=payload ) if response.status_code != 200: raise Exception(f"DeepSeek API 请求失败：{response.status_code} {response.text}") answer = response.json()["choices"][0]["message"]["content"] return { "answer": answer, "model": "deepseek-chat", "elapsed": 1.2, "cost": 0.002 } elif model == "text-only": # 本地简单文本模拟回复 answer = f"[text-only模式] 你的输入是：{prompt}" return { "answer": answer, "model": "text-only", "elapsed": 0, "cost": 0 } except Exception as e: print(f"❌ {model} 调用失败：{e}") continue print("="*30) return { "answer": "❌ 所有模型均调用失败，请检查网络或API Key。", "model": "none", "elapsed": 0, "cost": 0 } def extract_text_from_image(pil_image): result = reader.readtext(np.array(pil_image), detail=0) return [line.strip() for line in result if len(line.strip()) > 5] def build_index(vectors): dim = len(vectors[0]) index = faiss.IndexFlatL2(dim) index.add(np.array(vectors).astype("float32")) return index def search_index(index, query_vec, top_k=5): D, I = index.search(np.array([query_vec]).astype("float32"), top_k) return I[0].tolist() def fetch_related_papers(query): return [ {"title": "Document Understanding with OCR and LLMs", "url": "https://arxiv.org/abs/2403.01234"}, {"title": "Multimodal Large Models on Visual Text", "url": "https://arxiv.org/abs/2402.05678"} ] def clean_text(text: str) -> str: # 去除多余空行、控制字符，简单纠正断句 text = re.sub(r'\n\s*\n', '\n', text) # 连续空行合并 text = text.replace('\r', '') return text.strip() def chunk_text(text: str, max_chunk_size=1000) -> list: paragraphs = text.split('\n') chunks = [] current_chunk = [] current_len = 0 for para in paragraphs: para_len = len(para) if current_len + para_len > max_chunk_size and current_chunk: chunks.append('\n'.join(current_chunk)) current_chunk = [para] current_len = para_len else: current_chunk.append(para) current_len += para_len if current_chunk: chunks.append('\n'.join(current_chunk)) return chunks def model_call_func(prompt: str) -> str: # 简易同步调用示例，只调用 Claude，实际可改成 safe_call_model 同步包装 system_prompt = ( "你是一名专业的文档分析师，具备深厚的行业背景知识，" "能够根据用户上传的文档内容，结合上下文，进行深入的分析、总结和解读。" "即使内容简略，也要尝试推断合理信息，给出高质量的答案。" "请避免只字面翻译，要结合行业理解。" ) headers = { "x-api-key": os.getenv("ANTHROPIC_API_KEY"), "content-type": "application/json", "anthropic-version": "2023-06-01" } messages = [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt} ] payload = { "model": "claude-3-sonnet-20240229", "max_tokens": 512, "temperature": 0.5, "messages": messages } response = requests.post("https://api.anthropic.com/v1/messages", headers=headers, json=payload) if response.status_code != 200: print(f"摘要调用失败：{response.status_code} {response.text}") return "" data = response.json() if "content" in data and isinstance(data["content"], list): return data["content"][0].get("text", "") return "" def summarize_chunks(chunks: list, model_call_func) -> str: summaries = [] for chunk in chunks: prompt = f"请对以下内容进行简明扼要的总结：\n\n{chunk}" result = model_call_func(prompt) summaries.append(result) return "\n".join(summaries) def handle_file_upload(file, session_id): if file is None: return "❌ 请上传文件", [], False ext = os.path.splitext(file.name)[1].lower() if ext in [".jpg", ".jpeg", ".png"]: img = Image.open(file.name).convert("RGB") chunks = extract_text_from_image(img) gallery = [img] elif ext == ".pdf": from PyPDF2 import PdfReader reader = PdfReader(file.name) text = "\n".join(page.extract_text() or "" for page in reader.pages) chunks = text.split("\n\n") gallery = [] elif ext == ".docx": import docx doc = docx.Document(file.name) text = "\n".join(p.text for p in doc.paragraphs if p.text.strip()) chunks = text.split("\n\n") gallery = [] elif ext == ".txt": with open(file.name, "r", encoding="utf-8") as f: text = f.read() chunks = text.split("\n\n") gallery = [] elif ext == ".md": with open(file.name, "r", encoding="utf-8") as f: text = f.read() chunks = text.split("\n\n") gallery = [] elif ext == ".csv": import pandas as pd df = pd.read_csv(file.name) text = df.to_string() chunks = text.split("\n\n") gallery = [] elif ext == ".xlsx": import pandas as pd df = pd.read_excel(file.name, sheet_name=None) text = "\n\n".join([df[sheet].to_string() for sheet in df]) chunks = text.split("\n\n") gallery = [] else: return "❌ 不支持的文件格式", [], False if not chunks: return "❌ 未提取到有效文本内容", gallery, False # 清洗文本合并所有块 raw_text = "\n\n".join(chunks) cleaned_text = clean_text(raw_text) # 分段拆分 chunk_list = chunk_text(cleaned_text, max_chunk_size=1500) # 生成摘要 summary_text = summarize_chunks(chunk_list, model_call_func) # 将摘要插入chunks首位，方便后续检索提升相关性 chunks.insert(0, "[摘要]\n" + summary_text) vectors = embed_texts(chunks) index = build_index(vectors) session_data[session_id] = {"chunks": chunks, "index": index, "history": []} return "✅ 识别成功，已建立索引", gallery, True def handle_user_query(query, session_id, model_name, chat_history, image_uploaded, task_mode="default"): if task_mode == "summary": query = "请对上传的内容进行摘要提炼。" elif task_mode == "translation": query = "请将上传的内容翻译成中文。" elif task_mode == "qa-extract": query = "请从上传内容中提取结构化的问答对。" # 文本模式优先跳过OCR索引 if model_name == "text-only": history = session_data.get(session_id, {}).get("history", []) prompt = f"""你是一个多模态文档理解助手，能够总结上传的文件、图像或纯文本内容。请根据以下用户输入的问题进行回答。 用户问题：{query} """ else: if image_uploaded and session_id in session_data: chunks = session_data[session_id]["chunks"] index = session_data[session_id]["index"] history = session_data[session_id]["history"] q_vec = embed_texts([query])[0] top_ids = search_index(index, q_vec) context = "\n".join([chunks[i] for i in top_ids]) MAX_PROMPT_LEN = 4000 if len(context) > MAX_PROMPT_LEN: context = context[:MAX_PROMPT_LEN] prompt = f"""你是一个文档理解助手，具备结构化内容分析能力。请基于以下文档内容尽可能准确地回答用户提出的问题。文档内容可能来自 OCR 图片、PDF、Word、Excel 等。 文档内容如下： {context} 用户问题：{query} """ else: history = session_data.get(session_id, {}).get("history", []) prompt = f"""你是一个多模态文档理解助手，能够总结上传的文件、图像或纯文本内容。请根据以下用户输入的问题进行回答。 用户问题：{query} """ model_list = [model_name] if model_name != "auto" else ["claude", "deepseek", "openai"] result = safe_call_model(prompt, session_id, model_list, history) papers = fetch_related_papers(query) paper_text = "\n\n📚 相关研究论文：\n" + "\n".join([f"- {p['title']}（{p['url']}）" for p in papers]) combined_answer = result["answer"] + paper_text history += [{"role": "user", "content": query}, {"role": "assistant", "content": combined_answer}] if session_id in session_data: session_data[session_id]["history"] = history else: session_data[session_id] = {"history": history, "chunks": [], "index": None} return history, f"模型：{result['model']} 耗时：{result['elapsed']:.2f}s 估算费用：${result['cost']:.6f}" with gr.Blocks() as demo: session_id = gr.State(f"session_{int(time.time())}_{np.random.randint(1000,9999)}") gr.Markdown("# 📷 图片问答 + 多模型文本对话系统") with gr.Row(): image_upload = gr.File(label="上传文件（支持图片 / PDF / Word / 文本 / Excel）", file_types=[".jpg", ".png", ".jpeg", ".pdf", ".docx", ".txt", ".md", ".csv", ".xlsx"]) upload_btn = gr.Button("识别并建立索引") upload_status = gr.Textbox(label="上传状态", interactive=False) gallery = gr.Gallery(label="上传图像") model_select = gr.Radio( ["auto", "claude", "deepseek", "openai", "text-only"], label="选择模型", value="auto" ) task_mode = gr.Radio( ["default", "summary", "translation", "qa-extract"], label="任务类型", value="default" ) chatbot = gr.Chatbot(label="聊天记录", height=400, type="messages") query_input = gr.Textbox(label="请输入问题") query_btn = gr.Button("发送查询") status_output = gr.Textbox(label="状态", interactive=False) image_uploaded = gr.State(False) upload_btn.click( fn=handle_file_upload, inputs=[image_upload, session_id], outputs=[upload_status, gallery, image_uploaded] ) query_btn.click( fn=handle_user_query, inputs=[query_input, session_id, model_select, chatbot, image_uploaded, task_mode], outputs=[chatbot, status_output] ) # 支持回车提交查询 query_input.submit( fn=handle_user_query, inputs=[query_input, session_id, model_select, chatbot, image_uploaded, task_mode], outputs=[chatbot, status_output] ) if __name__ == "__main__": demo.launch()