import torch from transformers import AutoTokenizer, AutoModelForTokenClassification, pipeline, AutoModel import gradio as gr import re import os 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 # 配置日志 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") # 使用SQLAlchemy的连接池来管理数据库连接 DATABASE_URL = "mysql+pymysql://user:password@host/dbname" # 请根据实际情况修改连接字符串 # 创建引擎(连接池) engine = create_engine(DATABASE_URL, pool_size=10, max_overflow=20, echo=True) # 创建session类 Session = sessionmaker(bind=engine) @contextmanager def get_db_connection(): """ 使用上下文管理器获取数据库连接 """ session = None try: session = Session() # 从连接池中获取一个连接 logging.info("✅ 数据库连接已建立") yield session # 使用session进行数据库操作 except Exception as e: logging.error(f"❌ 数据库操作时发生错误: {e}") if session: session.rollback() # 回滚事务 finally: if session: try: session.commit() # 提交事务 logging.info("✅ 数据库事务已提交") except Exception as e: logging.error(f"❌ 提交事务时发生错误: {e}") finally: session.close() # 关闭会话,释放连接 logging.info("✅ 数据库连接已关闭") def save_to_db(table, data): """ 将数据保存到数据库 :param table: 表名 :param data: 数据字典 """ try: valid_tables = ["entities", "relations"] # 只允许保存到这些表 if table not in valid_tables: raise ValueError(f"Invalid table: {table}") with get_db_connection() as conn: if conn: # 这里的操作假设使用了ORM模型来处理插入,实际根据你数据库的表结构来调整 table_model = get_table_model(table) # 假设你有一个方法来根据表名获得ORM模型 new_record = table_model(**data) conn.add(new_record) conn.commit() # 提交事务 except Exception as e: logging.error(f"❌ 保存数据时发生错误: {e}") return False return True def get_table_model(table_name): """ 根据表名获取ORM模型(这里假设你有一个映射到数据库表的模型) :param table_name: 表名 :return: 对应的ORM模型 """ if table_name == "entities": from models import Entity # 假设你已经定义了ORM模型 return Entity elif table_name == "relations": from models import Relation # 假设你已经定义了ORM模型 return Relation else: raise ValueError(f"Unknown table: {table_name}") # ======================== 模型加载 ======================== NER_MODEL_NAME = "uer/roberta-base-finetuned-cluener2020-chinese" @lru_cache(maxsize=1) def get_ner_pipeline(): tokenizer = AutoTokenizer.from_pretrained(NER_MODEL_NAME) model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_NAME) return pipeline( "ner", model=model, tokenizer=tokenizer, aggregation_strategy="first" ) @lru_cache(maxsize=1) def get_re_pipeline(): return pipeline( "text2text-generation", model=NER_MODEL_NAME, tokenizer=NER_MODEL_NAME, max_length=512, device=0 if torch.cuda.is_available() else -1 ) # chatglm_model, chatglm_tokenizer = None, None # use_chatglm = False # try: # chatglm_model_name = "THUDM/chatglm-6b-int4" # chatglm_tokenizer = AutoTokenizer.from_pretrained(chatglm_model_name, trust_remote_code=True) # chatglm_model = AutoModel.from_pretrained( # chatglm_model_name, # trust_remote_code=True, # device_map="cpu", # torch_dtype=torch.float32 # ).eval() # use_chatglm = True # print("✅ 4-bit量化版ChatGLM加载成功") # except Exception as e: # print(f"❌ ChatGLM加载失败: {e}") # ======================== 知识图谱结构 ======================== knowledge_graph = {"entities": set(), "relations": set()} # 优化知识图谱更新函数,增加全局变量更新 def update_knowledge_graph(entities, relations): """ 更新知识图谱并保存到数据库 """ global knowledge_graph # 明确声明使用全局变量 # 保存实体 for e in entities: if isinstance(e, dict) and 'text' in e and 'type' in e: save_to_db('entities', { 'text': e['text'], 'type': e['type'], 'start_pos': e.get('start', -1), 'end_pos': e.get('end', -1), 'source': 'user_input' }) knowledge_graph["entities"].add((e['text'], e['type'])) # 保存关系 for r in relations: if isinstance(r, dict) and all(k in r for k in ("head", "tail", "relation")): save_to_db('relations', { 'head_entity': r['head'], 'tail_entity': r['tail'], 'relation_type': r['relation'], 'source_text': '' # 可添加原文关联 }) knowledge_graph["relations"].add((r['head'], r['tail'], r['relation'])) # 优化知识图谱文本格式生成函数,增加排序和去重 def visualize_kg_text(): """ 生成知识图谱的文本格式 """ nodes = sorted(set([f"{ent[0]} ({ent[1]})" for ent in knowledge_graph["entities"]])) edges = sorted(set([f"{h} --[{r}]-> {t}" for h, t, r in knowledge_graph["relations"]])) return "\n".join(["📌 实体:"] + nodes + ["", "📎 关系:"] + edges) # 优化知识图谱可视化函数,动态生成HTML文件名,避免覆盖 def visualize_kg_interactive(entities, relations): """ 生成交互式的知识图谱可视化 """ # 创建一个新的网络图 net = Network(height="700px", width="100%", bgcolor="#ffffff", font_color="black") # 定义实体类型颜色 entity_colors = { 'PER': '#FF6B6B', # 人物-红色 'ORG': '#4ECDC4', # 组织-青色 'LOC': '#45B7D1', # 地点-蓝色 'TIME': '#96CEB4', # 时间-绿色 'TITLE': '#D4A5A5' # 职位-灰色 } # 添加实体节点 for entity in entities: node_color = entity_colors.get(entity['type'], '#D3D3D3') # 默认灰色 net.add_node(entity['text'], label=f"{entity['text']} ({entity['type']})", color=node_color, title=f"类型: {entity['type']}") # 添加关系边 for relation in relations: net.add_edge(relation['head'], relation['tail'], label=relation['relation'], arrows='to') # 设置物理布局 net.set_options(''' var options = { "physics": { "forceAtlas2Based": { "gravitationalConstant": -50, "centralGravity": 0.01, "springLength": 100, "springConstant": 0.08 }, "maxVelocity": 50, "solver": "forceAtlas2Based", "timestep": 0.35, "stabilization": {"iterations": 150} } } ''') # 动态生成HTML文件名 timestamp = int(time.time()) html_path = f"knowledge_graph_{timestamp}.html" net.save_graph(html_path) return html_path # ======================== 实体识别(NER) ======================== def merge_adjacent_entities(entities): if not entities: return entities merged = [entities[0]] for entity in entities[1:]: last = merged[-1] # 合并相邻的同类型实体 if (entity["type"] == last["type"] and entity["start"] == last["end"]): last["text"] += entity["text"] last["end"] = entity["end"] else: merged.append(entity) return merged def ner(text, model_type="bert"): start_time = time.time() # 如果使用的是 ChatGLM 模型,执行 ChatGLM 的NER if model_type == "chatglm" and use_chatglm: try: prompt = f"""请从以下文本中识别所有实体,严格按照JSON列表格式返回,每个实体包含text、type、start、end字段: 示例:[{{"text": "北京", "type": "LOC", "start": 0, "end": 2}}] 文本:{text}""" response = chatglm_model.chat(chatglm_tokenizer, prompt, temperature=0.1) if isinstance(response, tuple): response = response[0] try: json_str = re.search(r'\[.*\]', response, re.DOTALL).group() entities = json.loads(json_str) valid_entities = [ent for ent in entities if all(k in ent for k in ("text", "type", "start", "end"))] return valid_entities, time.time() - start_time except Exception as e: print(f"JSON解析失败: {e}") return [], time.time() - start_time except Exception as e: print(f"ChatGLM调用失败: {e}") return [], time.time() - start_time # 使用BERT NER text_chunks = [text[i:i + 510] for i in range(0, len(text), 510)] # 安全分段 raw_results = [] # 获取NER pipeline ner_pipeline = get_ner_pipeline() # 使用缓存的pipeline for idx, chunk in enumerate(text_chunks): chunk_results = ner_pipeline(chunk) # 使用获取的pipeline for r in chunk_results: r["start"] += idx * 510 r["end"] += idx * 510 raw_results.extend(chunk_results) entities = [{ "text": r['word'].replace(' ', ''), "start": r['start'], "end": r['end'], "type": LABEL_MAPPING.get(r.get('entity_group') or r.get('entity'), r.get('entity_group') or r.get('entity')) } for r in raw_results] entities = merge_adjacent_entities(entities) return entities, time.time() - start_time # ------------------ 实体类型标准化 ------------------ LABEL_MAPPING = { "address": "LOC", "company": "ORG", "name": "PER", "organization": "ORG", "position": "TITLE", "government": "ORG", "scene": "LOC", "book": "WORK", "movie": "WORK", "game": "WORK" } # 提取实体 entities, processing_time = ner("Google in New York met Alice") # 标准化实体类型 for e in entities: e["type"] = LABEL_MAPPING.get(e.get("type"), e.get("type")) # 打印标准化后的实体 print(f"[DEBUG] 标准化后实体列表: {[{'text': e['text'], 'type': e['type']} for e in entities]}") # 打印处理时间 print(f"处理时间: {processing_time:.2f}秒") # ======================== 关系抽取(RE) ======================== @lru_cache(maxsize=1) def get_re_pipeline(): tokenizer = AutoTokenizer.from_pretrained(NER_MODEL_NAME) model = AutoModelForTokenClassification.from_pretrained(NER_MODEL_NAME) return pipeline( "ner", # 使用NER pipeline model=model, tokenizer=tokenizer, aggregation_strategy="first" ) def re_extract(entities, text, use_bert_model=True): if not entities or not text: return [], 0 start_time = time.time() try: # 使用规则匹配关系 relations = [] # 定义关系关键词和对应的实体类型约束 relation_rules = { "位于": { "keywords": ["位于", "在", "坐落于"], "valid_types": { "head": ["ORG", "PER", "LOC"], "tail": ["LOC"] } }, "属于": { "keywords": ["属于", "是", "为"], "valid_types": { "head": ["ORG", "PER"], "tail": ["ORG", "LOC"] } }, "任职于": { "keywords": ["任职于", "就职于", "工作于"], "valid_types": { "head": ["PER"], "tail": ["ORG"] } } } # 预处理实体,去除重复和部分匹配 processed_entities = [] for e in entities: # 检查是否与已有实体重叠 is_subset = False for pe in processed_entities: if e["text"] in pe["text"] and e["text"] != pe["text"]: is_subset = True break if not is_subset: processed_entities.append(e) # 遍历文本中的每个句子 sentences = re.split('[。!?.!?]', text) for sentence in sentences: if not sentence.strip(): continue # 获取当前句子中的实体 sentence_entities = [e for e in processed_entities if e["text"] in sentence] # 检查每个关系类型 for rel_type, rule in relation_rules.items(): for keyword in rule["keywords"]: if keyword in sentence: # 在句子中查找符合类型约束的实体对 for i, ent1 in enumerate(sentence_entities): for j, ent2 in enumerate(sentence_entities): if i != j: # 避免自循环 # 检查实体类型是否符合规则 if (ent1["type"] in rule["valid_types"]["head"] and ent2["type"] in rule["valid_types"]["tail"]): # 检查实体在句子中的位置关系 if sentence.find(ent1["text"]) < sentence.find(ent2["text"]): relations.append({ "head": ent1["text"], "tail": ent2["text"], "relation": rel_type }) # 去重 unique_relations = [] seen = set() for rel in relations: rel_key = (rel["head"], rel["tail"], rel["relation"]) if rel_key not in seen: seen.add(rel_key) unique_relations.append(rel) return unique_relations, time.time() - start_time except Exception as e: logging.error(f"关系抽取失败: {e}") return [], time.time() - start_time # ======================== 文本分析主流程 ======================== def create_knowledge_graph(entities, relations): """ 创建知识图谱可视化(文本格式) """ # 设置实体类型的颜色映射 entity_colors = { 'PER': '🔴', # 人物-红色 'ORG': '🔵', # 组织-蓝色 'LOC': '🟢', # 地点-绿色 'TIME': '🟡', # 时间-黄色 'TITLE': '🟣' # 职位-紫色 } # 生成实体列表 entity_list = [] for entity in entities: emoji = entity_colors.get(entity['type'], '⚪') entity_list.append(f"{emoji} {entity['text']} ({entity['type']})") # 生成关系列表 relation_list = [] for relation in relations: relation_list.append(f"{relation['head']} --[{relation['relation']}]--> {relation['tail']}") # 生成HTML内容 html_content = f"""