triple_extraction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "from sklearn.metrics.pairwise import cosine_similarity\n",
    "from joblib import Memory\n",
    "import json\n",
    "import openai\n",
    "import os\n",
    "import sys\n",
    "from tenacity import retry, stop_after_attempt, wait_random_exponential\n",
    "import time\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "import torch\n",
    "from langchain import HuggingFacePipeline\n",
    "from langchain.vectorstores import Chroma\n",
    "from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter\n",
    "from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM, TextStreamer, GPTQConfig\n",
    "from kg_rag.config_loader import *\n",
    "import ast\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pathlib import Path\n",
    "import sys\n",
    "\n",
    "# 获取当前工作目录的绝对路径\n",
    "current_dir = Path.cwd()\n",
    "\n",
    "# 将配置文件的路径添加到系统路径中\n",
    "config_path = os.path.join(current_dir, \"key_config.py\")\n",
    "if config_path not in sys.path:\n",
    "    sys.path.append(config_path)\n",
    "\n",
    "# # 导入 API 密钥\n",
    "from key_config import OPENAI_KEY\n",
    "from key_config import DASHSCOPE_KEY\n",
    "\n",
    "# # 现在可以使用 API_KEY 变量了\n",
    "print(OPENAI_KEY)\n",
    "print(DASHSCOPE_KEY)\n",
    "\n",
    "client = openai.OpenAI(api_key=OPENAI_KEY)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#爬取论文一级消息链接\n",
    "import urllib.request\n",
    "import time\n",
    "import feedparser\n",
    "from http import HTTPStatus\n",
    "import dashscope\n",
    "\n",
    "# Base api query url\n",
    "base_url = 'http://export.arxiv.org/api/query?';\n",
    "\n",
    "# Search parameters\n",
    "search_query = 'all:llm' # search for electron in all fields\n",
    "start = 0                       # start at the first result\n",
    "total_results = 20              # want 20 total results\n",
    "results_per_iteration = 20       # 5 results at a time\n",
    "wait_time = 1                   # number of seconds to wait beetween calls\n",
    "\n",
    "\n",
    "\n",
    "# entities_prompt = '''Now that you are an entity recognition expert, you are now given a text, you need to identify all the entities of the text,and return Entities in the following json format: \\nEntities:<List of extracted entities>\\nPlease return only the json format of the entity and nothing else'''\n",
    "\n",
    "abstract_prompt = '''Now that you are now an expert in summarizing summaries, you are now given a abstract of an article, please extract from it which field it belongs to (like continuous learning, large language models, reinforcement learning, etc.), problems, methods and effects, and return it in the following json format \\n\\{\"field\":, \"problems\":, \"methods\":, \"effects\":\\}\\nPlease return only the json and nothing else'''\n",
    "\n",
    "\n",
    "def fetch_GPT_response(instruction, base_prompt, chat_model_id, temperature=0):\n",
    "    print('Calling OpenAI...')\n",
    "    response = client.chat.completions.create(\n",
    "        model=chat_model_id,\n",
    "        messages=[\n",
    "            # {\"role\": \"system\", \"content\": base_prompt},\n",
    "            {\"role\": \"user\", \"content\": instruction}\n",
    "        ],\n",
    "        temperature=temperature,\n",
    "    )\n",
    "    return response.choices[0].message.content.strip()\n",
    "\n",
    "def get_GPT_response(instruction, base_prompt, chat_model_id, temperature=0):\n",
    "    return fetch_GPT_response(instruction, base_prompt, chat_model_id, temperature)\n",
    "\n",
    "dashscope.api_key = DASHSCOPE_KEY\n",
    "\n",
    "def call_with_messages(system_prompts, question):\n",
    "    messages = [{'role': 'system', 'content': system_prompts},\n",
    "                {'role': 'user', 'content': question}]\n",
    "\n",
    "    response = dashscope.Generation.call(\n",
    "        dashscope.Generation.Models.qwen_72b_chat,\n",
    "        messages=messages,\n",
    "        result_format='message',  # set the result to be \"message\" format.\n",
    "    )\n",
    "    if response.status_code == HTTPStatus.OK:\n",
    "        print(response)\n",
    "    else:\n",
    "        print('Request id: %s, Status code: %s, error code: %s, error message: %s' % (\n",
    "            response.request_id, response.status_code,\n",
    "            response.code, response.message\n",
    "        ))\n",
    "        \n",
    "# call_with_messages()\n",
    "\n",
    "def abstract_extractor(text):\n",
    "    # chat_model_id, chat_deployment_id = get_gpt35()\n",
    "    prompt_updated = abstract_prompt+ \"\\n\" + \"Text : \" + text\n",
    "    print('prompt_updated:',prompt_updated)\n",
    "\n",
    "    resp = get_GPT_response(prompt_updated, abstract_prompt, 'gpt-4', temperature=0)\n",
    "    print('resp:',resp)\n",
    "    try:\n",
    "        abstract_dict = json.loads(resp)\n",
    "        print(abstract_dict)\n",
    "        # return entity_dict[\"Diseases\"]\n",
    "        return abstract_dict\n",
    "    except:\n",
    "        return None\n",
    "\n",
    "print ('Searching arXiv for %s' % search_query)\n",
    "\n",
    "import json\n",
    "\n",
    "# 初始化数据\n",
    "title = \"A Comprehensive Survey on Code Large Language Models\"\n",
    "author = \"Jane Doe, John Smith\"\n",
    "published = \"2024-02-28\"\n",
    "\n",
    "resp = {\n",
    "    'field': 'Large Language Models, Software Engineering',\n",
    "    'problems': 'Lack of systematic investigation into Code LLMs and their performance, frequent updates of Code LLMs influenced by base LLMs',\n",
    "    'methods': 'Comprehensive survey and analysis of Code LLMs, categorization of Code LLMs based on their publishers, investigation of performance differences between general LLMs and Code LLMs, maintenance of performance of LLMs across multiple mainstream benchmarks',\n",
    "    'effects': 'Assists developers of Code LLMs in choosing base models for the development of more advanced LLMs, provides insights for practitioners to better understand key improvement directions for Code LLMs'\n",
    "}\n",
    "\n",
    "knowledge_graph = {\n",
    "    \"edges\": [],\n",
    "    \"entities\": {}\n",
    "}\n",
    "\n",
    "# 初始构建知识图谱函数\n",
    "def build_knowledge_graph(title, author, published, resp):\n",
    "    print('resp:',resp)\n",
    "    # 如果作者和发布日期还未作为独立的实体加入，则加入它们\n",
    "    if author not in knowledge_graph[\"entities\"]:\n",
    "        knowledge_graph[\"entities\"][author] = {}\n",
    "        knowledge_graph[\"edges\"].append([title, \"author\", author])\n",
    "    \n",
    "    if published not in knowledge_graph[\"entities\"]:\n",
    "        knowledge_graph[\"entities\"][published] = {}\n",
    "        knowledge_graph[\"edges\"].append([title, \"published\", published])\n",
    "        \n",
    "    if resp['field'] not in knowledge_graph['entities']:\n",
    "        knowledge_graph[\"entities\"][resp['field']] = {}\n",
    "        \n",
    "    # print(resp)\n",
    "        \n",
    "    knowledge_graph['entities'][title] = resp\n",
    "    \n",
    "\n",
    "\n",
    "for i in range(start,total_results,results_per_iteration):\n",
    "    \n",
    "    print (\"Results %i - %i\" % (i,i+results_per_iteration))\n",
    "    \n",
    "    query = 'search_query=%s&start=%i&max_results=%i' % (search_query,\n",
    "                                                         i,\n",
    "                                                        results_per_iteration)\n",
    "\n",
    "    # perform a GET request using the base_url and query\n",
    "    response = urllib.request.urlopen(base_url+query).read()\n",
    "\n",
    "    # parse the response using feedparser\n",
    "    feed = feedparser.parse(response)\n",
    "\n",
    "    # Run through each entry, and print out information\n",
    "    for entry in feed.entries:\n",
    "        print ('arxiv-id: %s' % entry.id.split('/abs/')[-1])\n",
    "        print ('Link:  %s' % entry.id)\n",
    "        print ('Title:  %s' % entry.title)\n",
    "        print ('Abstract: %s' %entry.summary.replace(\"\\n\",\"\"))\n",
    "        # feedparser v4.1 only grabs the first author\n",
    "        print ('First Author:  %s' % entry.author)\n",
    "        print ('Publish_time: %s' % entry.published)\n",
    "        print(\"=================================\")\n",
    "        # 使用初始函数构建知识图谱\n",
    "\n",
    "        resp = abstract_extractor(entry.summary.replace(\"\\n\",\"\"))\n",
    "        build_knowledge_graph(entry.title, entry.author, entry.published, resp)\n",
    "    \n",
    "    # Remember to play nice and sleep a bit before you call\n",
    "    # the api again!\n",
    "    print ('Sleeping for %i seconds' % wait_time )\n",
    "    time.sleep(wait_time)\n",
    "\n",
    "\n",
    "file_name = 'knowledge_graph_with_dynamic_entities.json'\n",
    "with open(file_name, 'w') as file:\n",
    "    json.dump(knowledge_graph, file, indent=4)\n",
    "\n",
    "print(f\"Knowledge graph saved to {file_name}.\")\n",
    "# #下载指定链接的pdf\n",
    "# from urllib import request\n",
    "\n",
    "# url = \"https://arxiv.org/pdf/2311.07989.pdf\"\n",
    "# request.urlretrieve(url,'1.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "# 初始化数据\n",
    "title = \"A Comprehensive Survey on Code Large Language Models\"\n",
    "author = \"Jane Doe, John Smith\"\n",
    "published = \"2024-02-28\"\n",
    "\n",
    "resp = {\n",
    "    'field': 'Large Language Models, Software Engineering',\n",
    "    'problems': 'Lack of systematic investigation into Code LLMs and their performance, frequent updates of Code LLMs influenced by base LLMs',\n",
    "    'methods': 'Comprehensive survey and analysis of Code LLMs, categorization of Code LLMs based on their publishers, investigation of performance differences between general LLMs and Code LLMs, maintenance of performance of LLMs across multiple mainstream benchmarks',\n",
    "    'effects': 'Assists developers of Code LLMs in choosing base models for the development of more advanced LLMs, provides insights for practitioners to better understand key improvement directions for Code LLMs'\n",
    "}\n",
    "\n",
    "KG = {\n",
    "    \"edges\": [],\n",
    "    \"entities\": {}\n",
    "}\n",
    "\n",
    "# 初始构建知识图谱函数\n",
    "def build_knowledge_graph(title, author, published, resp):\n",
    "    # 如果作者和发布日期还未作为独立的实体加入，则加入它们\n",
    "    if author not in knowledge_graph[\"entities\"]:\n",
    "        KG[\"entities\"][author] = {}\n",
    "        KG[\"edges\"].append([title, \"author\", author])\n",
    "    \n",
    "    if published not in knowledge_graph[\"entities\"]:\n",
    "        KG[\"entities\"][published] = {}\n",
    "        KG[\"edges\"].append([title, \"published\", published])\n",
    "        \n",
    "    if resp['field'] not in knowledge_graph['entities']:\n",
    "        KG[\"entities\"][resp['field']] = {}\n",
    "        \n",
    "    # print(resp)\n",
    "        \n",
    "    KG['entities'][title] = resp\n",
    "    \n",
    "\n",
    "# 使用初始函数构建知识图谱\n",
    "build_knowledge_graph(title, author, published, resp)\n",
    "\n",
    "# # 假设有新的实体数据传入\n",
    "# new_entity_name = \"NewEntity\"\n",
    "# new_entity_info = {\"attribute1\": \"value1\", \"attribute2\": \"value2\"}\n",
    "\n",
    "# # 向知识图谱添加新实体\n",
    "# add_entity_to_graph(knowledge_graph, new_entity_name, new_entity_info)\n",
    "\n",
    "# file_name = 'knowledge_graph_with_dynamic_entities.json'\n",
    "# with open(file_name, 'w') as file:\n",
    "#     json.dump(knowledge_graph, file, indent=4)\n",
    "\n",
    "print(json.dumps(KG, indent=4))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "knowledge_graph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# entities_prompt = '''Now that you are an entity recognition expert, you are now given a text, you need to identify all the entities of the text,and return Entities in the following json format: \\nEntities:<List of extracted entities>\\nPlease return only the json format of the entity and nothing else'''\n",
    "\n",
    "abstract_prompt = '''Now that you are now an expert in summarizing summaries, you are now given a abstract of an article, please extract from it which field it belongs to (like continuous learning, large language models, reinforcement learning, etc.), problems, methods and effects, and return it in the following json format \\n\\{\"field\":, \"problems\":, \"methods\":, \"effects\":\\}\\nPlease return only the json and nothing else'''\n",
    "\n",
    "\n",
    "def fetch_GPT_response(instruction, base_prompt, chat_model_id, temperature=0):\n",
    "    print('Calling OpenAI...')\n",
    "    response = client.chat.completions.create(\n",
    "        model=chat_model_id,\n",
    "        messages=[\n",
    "            # {\"role\": \"system\", \"content\": base_prompt},\n",
    "            {\"role\": \"user\", \"content\": instruction}\n",
    "        ],\n",
    "        temperature=temperature,\n",
    "    )\n",
    "    return response.choices[0].message.content.strip()\n",
    "\n",
    "def get_GPT_response(instruction, base_prompt, chat_model_id, temperature=0):\n",
    "    return fetch_GPT_response(instruction, base_prompt, chat_model_id, temperature)\n",
    "\n",
    "def abstract_extractor(text):\n",
    "    # chat_model_id, chat_deployment_id = get_gpt35()\n",
    "    prompt_updated = abstract_prompt+ \"\\n\" + \"Text : \" + text\n",
    "    print('prompt_updated:',prompt_updated)\n",
    "\n",
    "    resp = get_GPT_response(prompt_updated, abstract_prompt, 'gpt-4', temperature=0)\n",
    "    print('resp:',resp)\n",
    "    abstract_dict = json.loads(resp)\n",
    "    print(abstract_dict)\n",
    "    return abstract_dict\n",
    "\n",
    "    \n",
    "\n",
    "abstract = 'General large language models (LLMs), represented by ChatGPT, havedemonstrated significant potential in tasks such as code generation in softwareengineering. This has led to the development of specialized LLMs for softwareengineering, known as Code LLMs. A considerable portion of Code LLMs is derivedfrom general LLMs through model fine-tuning. As a result, Code LLMs are oftenupdated frequently and their performance can be influenced by the base LLMs.However, there is currently a lack of systematic investigation into Code LLMsand their performance. In this study, we conduct a comprehensive survey andanalysis of the types of Code LLMs and their differences in performancecompared to general LLMs. We aim to address three questions: (1) What LLMs arespecifically designed for software engineering tasks, and what is therelationship between these Code LLMs? (2) Do Code LLMs really outperformgeneral LLMs in software engineering tasks? (3) Which LLMs are more proficientin different software engineering tasks? To answer these questions, we firstcollect relevant literature and work from five major databases and open-sourcecommunities, resulting in 134 works for analysis. Next, we categorize the CodeLLMs based on their publishers and examine their relationships with generalLLMs and among themselves. Furthermore, we investigate the performancedifferences between general LLMs and Code LLMs in various software engineeringtasks to demonstrate the impact of base models and Code LLMs. Finally, wecomprehensively maintained the performance of LLMs across multiple mainstreambenchmarks to identify the best-performing LLMs for each software engineeringtask. Our research not only assists developers of Code LLMs in choosing basemodels for the development of more advanced LLMs but also provides insights forpractitioners to better understand key improvement directions for Code LLMs.'\n",
    "resp = abstract_extractor(abstract)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(resp)\n",
    "\n",
    "\n",
    "for key, value in resp.items():\n",
    "    print(f\"{key}:\")\n",
    "    for item in value.split(', '):\n",
    "        print(f\"    {item}\")\n",
    "    print(\"\")  # Adding an extra newline for better separation between sections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# 设置环境变量\n",
    "os.environ['http_proxy'] = 'http://gpu013:7890'\n",
    "os.environ['https_proxy'] = 'http://gpu013:7890' # 如果你也需要设置HTTPS代理\n",
    "\n",
    "# 验证环境变量是否设置成功\n",
    "print(os.environ['http_proxy'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file = config_data['GPT_CONFIG_FILE']\n",
    "load_dotenv(config_file)\n",
    "\n",
    "system_prompts[\"DISEASE_ENTITY_EXTRACTION\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "relation_prompt = '''You are now a relationship extraction expert, now given a series of entities, and the text related to this entity, you need to go through all the entities in pairs, if there is a relationship between the two entities, this relationship may not be inferred directly from the text, please return all their relationships in the following json format.\n",
    "{'triple':list of 'entities-relation-entities'},\n",
    "please return only the json result and nothing else.'''\n",
    "\n",
    "relation_prompt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "# entities_prompt = '''Now that you are an entity recognition expert, you are now given a text, you need to identify all the entities of the text,and return Entities in the following json format: \\nEntities:<List of extracted entities>\\nPlease return only the json format of the entity and nothing else'''\n",
    "\n",
    "entities_pro_prompt = '''Now that you are an entity recognition expert, you are now given a text, you need to identify all the entities of the text,and return Entities in the following json format: \\nEntities:<List of extracted entities>\\nPlease report only Policy. Do not report any other entities like Genes, Proteins, Enzymes etc.'''\n",
    "\n",
    "\n",
    "def fetch_GPT_response(instruction, base_prompt, chat_model_id, temperature=0):\n",
    "    print('Calling OpenAI...')\n",
    "    response = client.chat.completions.create(\n",
    "        model=chat_model_id,\n",
    "        messages=[\n",
    "            # {\"role\": \"system\", \"content\": base_prompt},\n",
    "            {\"role\": \"user\", \"content\": instruction}\n",
    "        ],\n",
    "        temperature=temperature,\n",
    "    )\n",
    "    return response.choices[0].message.content.strip()\n",
    "\n",
    "def get_GPT_response(instruction, base_prompt, chat_model_id, temperature=0):\n",
    "    return fetch_GPT_response(instruction, base_prompt, chat_model_id, temperature)\n",
    "\n",
    "def entity_extractor(text):\n",
    "    # chat_model_id, chat_deployment_id = get_gpt35()\n",
    "    prompt_updated = entities_pro_prompt+ \"\\n\" + \"Text : \" + text\n",
    "    print('prompt_updated:',prompt_updated)\n",
    "\n",
    "    resp = get_GPT_response(prompt_updated, entities_pro_prompt, 'gpt-4', temperature=0)\n",
    "    print('resp:',resp)\n",
    "    try:\n",
    "        entities_dict = json.loads(resp)\n",
    "        print(entities_dict)\n",
    "        # return entity_dict[\"Diseases\"]\n",
    "    except:\n",
    "        return None\n",
    "    \n",
    "    \n",
    "text = '''从国家整体层面看，2020年政务数据的相关政策上有两大热点。\n",
    "　　一是数据生产要素这一提法在政策层面得到了进一步明确，成为社会各界关注的热点。继《促进大数据发展行动纲要》、中央政治局第二次集体学习、十九届四中全会《决定》中提到数据要素之后，2020年3月30日，《中共中央国务院关于构建更加完善的要素市场化配置体制机制的意见》出台，这是中央关于要素市场化配置的第一份文件，明确提出“加快培育数据要素市场”，并强调要推进政府数据开放共享、提升社会数据资源价值、加强数据资源整合和安全保护的具体要求。\n",
    "　　二是高度关注数据安全和隐私保护。2020年，我国公布了《中华人民共和国数据安全法（草案）》、《中华人民共和国个人信息保护法（草案）》，并发起了《全球数据安全倡议》，旨在明确数据安全法律责任，完善监管体系，保障国家安全、公民个人隐私权益和社会安全稳定。文件中，对政务数据的安全与开放也提出了明确的要求。'''\n",
    "print(text)\n",
    "entity_extractor(text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# relation_prompt = '''You are now a relationship extraction expert, now given a series of entities, and the text related to this entity, you need to go through all the entities in pairs, if there is a relationship between the two entities, this relationship may not be inferred directly from the text, please return all their relationships in the following json format.\n",
    "# {'triple':list of 'entities-relation-entities'},\n",
    "# # please return only the json result and nothing else.'''\n",
    "\n",
    "relation_prompt = '''You are now a relationship extraction expert, now given a series of entities, and the text related to this entity, please return all their relationships in the following json format.\n",
    "{'triple':list of 'entities-relation-entities'},\n",
    "please return only the json result and nothing else.'''\n",
    "\n",
    "def triple_extractor(entities, text):\n",
    "    # entities:'e1,e2,e3……'\n",
    "    prompt_updated = relation_prompt+ \"\\n\" +\"entities:\"+ entities + \"\\n\" + \"Text : \" + text\n",
    "    print('prompt_updated:',prompt_updated)\n",
    "\n",
    "    resp = get_GPT_response(prompt_updated, relation_prompt, 'gpt-4', temperature=0)\n",
    "    print('resp:',resp)\n",
    "    try:\n",
    "        triple_dict = json.loads(resp)\n",
    "        print(triple_dict)\n",
    "        # return entity_dict[\"Diseases\"]\n",
    "    except:\n",
    "        return None\n",
    "    \n",
    "text = ''''小陈是小李的父亲，小李是小刚的父亲'''\n",
    "print(text)\n",
    "triple_extractor('小陈, 小李, 小刚', text)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "KG_RAG",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}