mcp_use.py

import sys
import os
sys.path.append(os.path.join(os.path.dirname(__file__), 'relation-extraction-master'))
import re
import torch
from gqlalchemy import Memgraph
from relation_extraction.hparams import hparams
from relation_extraction.model import SentenceRE
from relation_extraction.data_utils import MyTokenizer, get_idx2tag, convert_pos_to_mask

# 云端Memgraph连接参数
MEMGRAPH_HOST = '18.159.132.161'
MEMGRAPH_PORT = 7687
MEMGRAPH_USERNAME = 'b300000.de@gmail.com'
MEMGRAPH_PASSWORD = '159951Tjk.'  # 请替换为你的真实密码
MEMGRAPH_ENCRYPTED = True

# 连接memgraph云数据库
def get_memgraph_conn():
    return Memgraph(
        MEMGRAPH_HOST,
        MEMGRAPH_PORT,
        MEMGRAPH_USERNAME,
        MEMGRAPH_PASSWORD,
        encrypted=MEMGRAPH_ENCRYPTED
    )

# 单句预测，返回三元组
class RelationPredictor:
    def __init__(self, hparams):
        self.device = hparams.device
        torch.manual_seed(hparams.seed)
        self.idx2tag = get_idx2tag(hparams.tagset_file)
        hparams.tagset_size = len(self.idx2tag)
        self.model = SentenceRE(hparams).to(self.device)
        self.model.load_state_dict(torch.load(hparams.model_file))
        self.model.eval()
        self.tokenizer = MyTokenizer(hparams.pretrained_model_path)

    def predict_one(self, text, entity1, entity2):
        match_obj1 = re.search(entity1, text)
        match_obj2 = re.search(entity2, text)
        if not (match_obj1 and match_obj2):
            return None
        e1_pos = match_obj1.span()
        e2_pos = match_obj2.span()
        item = {
            'h': {'name': entity1, 'pos': e1_pos},
            't': {'name': entity2, 'pos': e2_pos},
            'text': text
        }
        tokens, pos_e1, pos_e2 = self.tokenizer.tokenize(item)
        encoded = self.tokenizer.bert_tokenizer.batch_encode_plus([(tokens, None)], return_tensors='pt')
        input_ids = encoded['input_ids'].to(self.device)
        token_type_ids = encoded['token_type_ids'].to(self.device)
        attention_mask = encoded['attention_mask'].to(self.device)
        e1_mask = torch.tensor([convert_pos_to_mask(pos_e1, max_len=attention_mask.shape[1])]).to(self.device)
        e2_mask = torch.tensor([convert_pos_to_mask(pos_e2, max_len=attention_mask.shape[1])]).to(self.device)
        with torch.no_grad():
            logits = self.model(input_ids, token_type_ids, attention_mask, e1_mask, e2_mask)[0]
            logits = logits.to(torch.device('cpu'))
        relation = self.idx2tag[logits.argmax(0).item()]
        return entity1, relation, entity2

# 写入memgraph
def insert_to_memgraph(memgraph, entity1, relation, entity2):
    memgraph.execute(
        "MERGE (a:Entity {name: $name1})",
        {"name1": entity1}
    )
    memgraph.execute(
        "MERGE (b:Entity {name: $name2})",
        {"name2": entity2}
    )
    memgraph.execute(
        f"MATCH (a:Entity {{name: $name1}}), (b:Entity {{name: $name2}}) MERGE (a)-[:{relation}]->(b)",
        {"name1": entity1, "name2": entity2}
    )

# 主流程

def main():
    memgraph = get_memgraph_conn()
    predictor = RelationPredictor(hparams)
    print("请输入句子和两个实体，识别关系并写入Memgraph。输入exit退出。")
    while True:
        text = input("输入中文句子：")
        if text.strip().lower() == 'exit':
            break
        entity1 = input("句子中的实体1：")
        if entity1.strip().lower() == 'exit':
            break
        entity2 = input("句子中的实体2：")
        if entity2.strip().lower() == 'exit':
            break
        result = predictor.predict_one(text, entity1, entity2)
        if result is None:
            print("实体未在句子中找到，请重试。")
            continue
        entity1, relation, entity2 = result
        insert_to_memgraph(memgraph, entity1, relation, entity2)
        print(f"已写入Memgraph：({entity1})-[:{relation}]->({entity2})")

if __name__ == '__main__':
    main()