Upload app.py

app.py

@@ -7,52 +7,7 @@ import json
 import chardet
 from sklearn.metrics import precision_score, recall_score, f1_score
 import time
-from functools import lru_cache
-from sqlalchemy import create_engine
-from sqlalchemy.orm import sessionmaker
-from contextlib import contextmanager
-import logging
-import networkx as nx
-from pyvis.network import Network
-import pandas as pd
-import matplotlib.pyplot as plt
-from gqlalchemy import Memgraph
-from mcp_use import RelationPredictor, insert_to_memgraph, get_memgraph_conn  # 引入mcp_use中的功能
-from relation_extraction.hparams import hparams  # 引入模型超参数
-
-# ======================== 数据库模块 ========================
-MEMGRAPH_HOST = '18.159.132.161'
-MEMGRAPH_PORT = 7687
-MEMGRAPH_USERNAME = '[email protected]'
-MEMGRAPH_PASSWORD = os.environ.get("MEMGRAPH_PASSWORD", "<YOUR MEMGRAPH PASSWORD HERE>")
-
-# 初始化 Memgraph 连接
-memgraph = get_memgraph_conn()
-
-# 初始化关系抽取模型
-relation_predictor = RelationPredictor(hparams)
-
-# ======================== 关系抽取功能整合 ========================
-def extract_and_save_relations(text, entity1, entity2):
-    """
-    使用 mcp_use.py 中的 RelationPredictor 提取关系，并保存到 Memgraph
-    """
-    try:
-        # 调用关系抽取模型
-        result = relation_predictor.predict_one(text, entity1, entity2)
-        if result is None:
-            return f"❌ 未找到实体 '{entity1}' 或 '{entity2}'"
-        
-        # 提取关系
-        entity1, relation, entity2 = result
-        
-        # 保存到 Memgraph
-        insert_to_memgraph(memgraph, entity1, relation, entity2)
-        return f"✅ 已写入 Memgraph：({entity1})-[:{relation}]->({entity2})"
-    except Exception as e:
-        logging.error(f"关系抽取失败: {e}")
-        return f"❌ 关系抽取失败: {e}"
-
+from functools import lru_cache  # 添加这行导入
 # ======================== 数据库模块 ========================
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
@@ -63,46 +18,6 @@ from pyvis.network import Network
 import pandas as pd
 import matplotlib.pyplot as plt
 
-
-from gqlalchemy import Memgraph
-import os
-MEMGRAPH_HOST = '18.159.132.161'
-MEMGRAPH_PORT = 7687
-MEMGRAPH_USERNAME = '[email protected]'
-MEMGRAPH_PASSWORD = os.environ.get("MEMGRAPH_PASSWORD", "<YOUR MEMGRAPH PASSWORD HERE>")
-
-def hello_memgraph():
-    """测试 Memgraph 数据库连接并进行健康检查"""
-    try:
-        connection = Memgraph(
-            host=os.environ["MEMGRAPH_HOST"],
-            port=int(os.environ["MEMGRAPH_PORT"]),
-            username=os.environ["MEMGRAPH_USERNAME"],
-            password=os.environ["MEMGRAPH_PASSWORD"],  # 强制从环境变量获取
-            encrypted=True,
-            ssl_verify=True,
-            ca_path="/etc/ssl/certs/memgraph.crt"
-        )
-        
-        # 健康检查查询
-        health = connection.execute_and_fetch("CALL mg.get('memgraph') YIELD value;")
-        health_status = next(health)["value"]["status"]
-        
-        # 创建测试节点
-        connection.execute(
-            'CREATE (n:ConnectionTest { message: "Hello Memgraph", ts: $ts })',
-            {"ts": datetime.now().isoformat()}
-        )
-        
-        return f"✅ 连接正常 | 状态: {health_status}"
-        
-    except Exception as e:
-        logging.error(f"连接失败: {str(e)}", exc_info=True)
-        return f"❌ 连接失败: {str(e)}"
-    finally:
-        connection.close()
-        
-
 # 配置日志
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 
@@ -593,55 +508,35 @@ def process_text(text, model_type="bert"):
     return ent_text, rel_text, kg_text, f"{total_duration:.2f} 秒"
 
 # ======================== 知识图谱可视化 ========================
-import matplotlib.pyplot as plt
-import networkx as nx
-import tempfile
-import os
-import logging
-from matplotlib import font_manager
-
-# 这个函数用于查找并验证中文字体路径
-def find_chinese_font():
-    # 尝试查找 Noto Sans CJK 字体
-    font_paths = [
-        "/usr/share/fonts/truetype/noto/NotoSansCJK-Regular.ttc",  # Noto CJK 字体
-        "/usr/share/fonts/truetype/wqy/wqy-microhei.ttc"         # 微软雅黑
-    ]
-    
-    for font_path in font_paths:
-        if os.path.exists(font_path):
-            logging.info(f"Found font at {font_path}")
-            return font_path
-    
-    logging.error("No Chinese font found!")
-    return None
-
 def generate_kg_image(entities, relations):
     """
-    中文知识图谱生成函数，支持自动匹配系统中的中文字体，避免中文显示为方框。
+    生成知识图谱的图片并保存到临时文件（Hugging Face适配版）
     """
     try:
-        # === 1. 确保使用合适的中文字体 ===
-        chinese_font = find_chinese_font()  # 调用查找字体函数
-        if chinese_font:
-            font_prop = font_manager.FontProperties(fname=chinese_font)
-            plt.rcParams['font.family'] = font_prop.get_name()
-        else:
-            # 如果字体路径未找到，使用默认字体（DejaVu Sans）
-            logging.warning("Using default font")
-            plt.rcParams['font.family'] = ['DejaVu Sans']
+        import tempfile
+        import matplotlib.pyplot as plt
+        import networkx as nx
+        import os
+
+        # === 1. 强制设置中文字体 ===
+        plt.rcParams['font.sans-serif'] = ['Noto Sans CJK SC']  # Hugging Face内置字体
+        plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
 
-        plt.rcParams['axes.unicode_minus'] = False
+        # === 2. 检查输入数据 ===
+        if not entities or not relations:
+            return None
 
-        # === 2. 创建图谱 ===
+        # === 3. 创建图谱 ===
         G = nx.DiGraph()
         entity_colors = {
-            'PER': '#FF6B6B',  # 人物-红色
-            'ORG': '#4ECDC4',  # 组织-青色
-            'LOC': '#45B7D1',  # 地点-蓝色
-            'TIME': '#96CEB4'  # 时间-绿色
+            'PER': '#FF6B6B',  # 红色
+            'ORG': '#4ECDC4',  # 青色
+            'LOC': '#45B7D1',  # 蓝色
+            'TIME': '#96CEB4', # 绿色
+            'TITLE': '#D4A5A5' # 灰色
         }
 
+        # 添加节点（实体）
         for entity in entities:
             G.add_node(
                 entity["text"],
@@ -649,6 +544,7 @@ def generate_kg_image(entities, relations):
                 color=entity_colors.get(entity['type'], '#D3D3D3')
             )
 
+        # 添加边（关系）
         for relation in relations:
             if relation["head"] in G.nodes and relation["tail"] in G.nodes:
                 G.add_edge(
@@ -657,15 +553,15 @@ def generate_kg_image(entities, relations):
                     label=relation["relation"]
                 )
 
-        # === 3. 绘图配置 ===
-        plt.figure(figsize=(12, 8), dpi=150)
-        pos = nx.spring_layout(G, k=0.7, seed=42)
+        # === 4. 绘图配置 ===
+        plt.figure(figsize=(12, 8), dpi=150)  # 降低DPI以节省内存
+        pos = nx.spring_layout(G, k=0.7, seed=42)  # 固定随机种子保证布局稳定
 
+        # 绘制节点和边
         nx.draw_networkx_nodes(
             G, pos,
             node_color=[G.nodes[n]['color'] for n in G.nodes],
-            node_size=800,
-            alpha=0.9
+            node_size=800
         )
         nx.draw_networkx_edges(
             G, pos,
@@ -675,44 +571,35 @@ def generate_kg_image(entities, relations):
             arrowsize=20
         )
 
-        node_labels = {n: G.nodes[n]['label'] for n in G.nodes}
+        # === 5. 绘制中文标签（关键修改点）===
         nx.draw_networkx_labels(
             G, pos,
-            labels=node_labels,
+            labels={n: G.nodes[n]['label'] for n in G.nodes},
             font_size=10,
-            font_family=font_prop.get_name() if chinese_font else 'SimHei',
-            font_weight='bold'
+            font_family='Noto Sans CJK SC'  # 显式指定字体
         )
-
-        edge_labels = nx.get_edge_attributes(G, 'label')
         nx.draw_networkx_edge_labels(
             G, pos,
-            edge_labels=edge_labels,
+            edge_labels=nx.get_edge_attributes(G, 'label'),
             font_size=8,
-            font_family=font_prop.get_name() if chinese_font else 'SimHei'
+            font_family='Noto Sans CJK SC'  # 显式指定字体
         )
 
         plt.axis('off')
-        plt.tight_layout()
-
-        # === 4. 保存图片 ===
-        temp_dir = tempfile.mkdtemp()  # 确保在 Docker 容器中有权限写入
-        output_path = os.path.join(temp_dir, "kg.png")
-
-        # 打印路径以方便调试
-        logging.info(f"Saving graph image to {output_path}")
-
-        plt.savefig(output_path, bbox_inches='tight', pad_inches=0.1)
+        
+        # === 6. 保存到临时文件 ===
+        temp_dir = tempfile.mkdtemp()
+        file_path = os.path.join(temp_dir, "kg.png")
+        plt.savefig(file_path, bbox_inches='tight')
         plt.close()
-
-        return output_path
-
+        
+        return file_path
+        
     except Exception as e:
-        logging.error(f"[ERROR] 图谱生成失败: {str(e)}")
+        logging.error(f"生成知识图谱图片失败: {str(e)}")
         return None
 
 
-# ======================== 文件处理 ========================
 def process_file(file, model_type="bert"):
     try:
         with open(file.name, 'rb') as f:
@@ -864,11 +751,13 @@ with gr.Blocks(css="""
 
     with gr.Tab("📄 文本分析"):
         input_text = gr.Textbox(lines=6, label="输入文本")
-        entity1 = gr.Textbox(label="实体1")
-        entity2 = gr.Textbox(label="实体2")
-        btn = gr.Button("提取关系并保存到 Memgraph")
-        output = gr.Textbox(label="结果")
-        btn.click(fn=extract_and_save_relations, inputs=[input_text, entity1, entity2], outputs=output)
+        model_type = gr.Radio(["bert", "chatglm"], value="bert", label="选择模型")
+        btn = gr.Button("开始分析")
+        out1 = gr.Textbox(label="识别实体")
+        out2 = gr.Textbox(label="识别关系")
+        out3 = gr.HTML(label="知识图谱")  # 使用HTML组件显示文本格式的知识图谱
+        out4 = gr.Textbox(label="耗时")
+        btn.click(fn=process_text, inputs=[input_text, model_type], outputs=[out1, out2, out3, out4])
 
     with gr.Tab("🗂 文件分析"):
         file_input = gr.File(file_types=[".txt", ".json"])
@@ -903,10 +792,4 @@ with gr.Blocks(css="""
         import_output = gr.Textbox(label="导入日志")
         import_btn.click(fn=lambda: import_dataset(dataset_path.value), outputs=import_output)
 
-        gr.Markdown("### 测试 Memgraph 数据库连接")
-        memgraph_btn = gr.Button("测试 Memgraph 连接")
-        memgraph_output = gr.Textbox(label="连接测试日志")
-        memgraph_btn.click(fn=hello_memgraph, outputs=memgraph_output)
-
-
 demo.launch(server_name="0.0.0.0", server_port=7860)