qwen_entity_kg_matching.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/cephfs/home/chendeyu/anaconda3/envs/KG_RAG/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "http://gpu013:7890\n"
     ]
    }
   ],
   "source": [
    "from http import HTTPStatus\n",
    "import dashscope\n",
    "import json\n",
    "\n",
    "import os\n",
    "\n",
    "from transformers import AutoTokenizer, AutoModelForSequenceClassification\n",
    "import torch\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'])"
   ]
  },
  {
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CUDA is available! Using GPU: NVIDIA TITAN Xp\n"
     ]
    }
   ],
   "source": [
    "from sentence_transformers import SentenceTransformer, util\n",
    "\n",
    "# 检查CUDA是否可用\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda\")\n",
    "    print(f\"CUDA is available! Using GPU: {torch.cuda.get_device_name(0)}\")\n",
    "else:\n",
    "    device = torch.device(\"cpu\")\n",
    "    print(\"CUDA is not available. Using CPU instead.\")\n",
    "\n",
    "# 加载模型\n",
    "model = SentenceTransformer('all-MiniLM-L6-v2')\n",
    "model = model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "dashscope.api_key = DASHSCOPE_KEY\n",
    "\n",
    "# system_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 {\"field\":, \"problems\":, \"methods\":, \"effects\":}Please return only the {...} not include json at the beginning and nothing else'''\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 {\"field\":\"...\", \"problems\":\"...\", \"methods\":\"...\", \"effects\":\"...\"}.'''\n",
    "\n",
    "\n",
    "def qwen_get_response(system_prompt, question):\n",
    "    messages = [{'role': 'system', 'content': system_prompt},\n",
    "                {'role': 'user', 'content': question}]\n",
    "\n",
    "    response = dashscope.Generation.call(\n",
    "        dashscope.Generation.Models.qwen_plus,\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",
    "        formatted_response = json.dumps(response, indent=4, ensure_ascii=False)\n",
    "        print(formatted_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",
    "    return formatted_response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "def qwen_72b_chat_entity_extractor(text):\n",
    "    entity_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,\n",
    "    and return Entities in the following json format: \\nEntities:<List of extracted entities>\\n.'''\n",
    "    resp = qwen_get_response(entity_prompt, text)\n",
    "    entity_dict = json.loads(resp)\n",
    "    print(entity_dict)\n",
    "    return entity_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "文本: 'self-discovery travel experience'\n",
      "匹配到 'Traveling helps you discover who you are and lets you experience the vastness of this world.', 相似度分数: 0.6171\n",
      "匹配到 'Reaching the summit of a mountain offers an exhilarating sense of achievement.', 相似度分数: 0.3315\n",
      "匹配到 'Immersing myself in the pages of a historical novel allows me to live in a different era.', 相似度分数: 0.3158\n",
      "\n",
      "================================================================================\n",
      "\n",
      "文本: 'mountain climbing achievement'\n",
      "匹配到 'Reaching the summit of a mountain offers an exhilarating sense of achievement.', 相似度分数: 0.7586\n",
      "匹配到 'Traveling helps you discover who you are and lets you experience the vastness of this world.', 相似度分数: 0.1685\n",
      "匹配到 'The joy of watching my garden grow after days of care is incredibly fulfilling.', 相似度分数: 0.1483\n",
      "\n",
      "================================================================================\n",
      "\n",
      "文本: 'morning forest tranquility'\n",
      "匹配到 'There is a special kind of calm in the early hours of the morning in a lush forest.', 相似度分数: 0.6990\n",
      "匹配到 'Immersing myself in the pages of a historical novel allows me to live in a different era.', 相似度分数: 0.2298\n",
      "匹配到 'The joy of watching my garden grow after days of care is incredibly fulfilling.', 相似度分数: 0.2129\n",
      "\n",
      "================================================================================\n",
      "\n",
      "文本: 'python coding creativity'\n",
      "匹配到 'Diving into learning Python has unlocked new creative potentials for me.', 相似度分数: 0.7134\n",
      "匹配到 'The act of preparing a meal with family is a source of joy and brings everyone closer.', 相似度分数: 0.2154\n",
      "匹配到 'Finding myself through the melodies I play on the piano is a profound experience.', 相似度分数: 0.2108\n",
      "\n",
      "================================================================================\n",
      "\n",
      "文本: 'family cooking bonding'\n",
      "匹配到 'The act of preparing a meal with family is a source of joy and brings everyone closer.', 相似度分数: 0.6349\n",
      "匹配到 'The joy of watching my garden grow after days of care is incredibly fulfilling.', 相似度分数: 0.1889\n",
      "匹配到 'Finding myself through the melodies I play on the piano is a profound experience.', 相似度分数: 0.0833\n",
      "\n",
      "================================================================================\n",
      "\n",
      "文本: 'historical fiction immersion'\n",
      "匹配到 'Immersing myself in the pages of a historical novel allows me to live in a different era.', 相似度分数: 0.6924\n",
      "匹配到 'Finding myself through the melodies I play on the piano is a profound experience.', 相似度分数: 0.2420\n",
      "匹配到 'Traveling helps you discover who you are and lets you experience the vastness of this world.', 相似度分数: 0.2178\n",
      "\n",
      "================================================================================\n",
      "\n",
      "文本: 'piano music self-expression'\n",
      "匹配到 'Finding myself through the melodies I play on the piano is a profound experience.', 相似度分数: 0.6157\n",
      "匹配到 'The act of preparing a meal with family is a source of joy and brings everyone closer.', 相似度分数: 0.1982\n",
      "匹配到 'Immersing myself in the pages of a historical novel allows me to live in a different era.', 相似度分数: 0.1762\n",
      "\n",
      "================================================================================\n",
      "\n",
      "文本: 'gardening nurturing fulfillment'\n",
      "匹配到 'The joy of watching my garden grow after days of care is incredibly fulfilling.', 相似度分数: 0.5938\n",
      "匹配到 'The act of preparing a meal with family is a source of joy and brings everyone closer.', 相似度分数: 0.3681\n",
      "匹配到 'There is a special kind of calm in the early hours of the morning in a lush forest.', 相似度分数: 0.2328\n",
      "\n",
      "================================================================================\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# 更新为包含意思相近的句子和文本\n",
    "json_data = {\n",
    "    \"entities\": {\n",
    "        \"Traveling helps you discover who you are and lets you experience the vastness of this world.\":{},  # Similar to entity1\n",
    "        \"Reaching the summit of a mountain offers an exhilarating sense of achievement.\":{},  # Similar to entity2\n",
    "        \"There is a special kind of calm in the early hours of the morning in a lush forest.\":{},  # Similar to entity3\n",
    "        \"Diving into learning Python has unlocked new creative potentials for me.\":{},  # Similar to entity4\n",
    "        \"The act of preparing a meal with family is a source of joy and brings everyone closer.\":{},  # Similar to entity5\n",
    "        \"Immersing myself in the pages of a historical novel allows me to live in a different era.\":{},  # Similar to entity6\n",
    "        \"Finding myself through the melodies I play on the piano is a profound experience.\":{},  # Similar to entity7\n",
    "        \"The joy of watching my garden grow after days of care is incredibly fulfilling.\":{}  # Similar to entity8\n",
    "    }\n",
    "}\n",
    "\n",
    "texts = [\n",
    "    \"self-discovery travel experience\",\n",
    "    \"mountain climbing achievement\",\n",
    "    \"morning forest tranquility\",\n",
    "    \"python coding creativity\",\n",
    "    \"family cooking bonding\",\n",
    "    \"historical fiction immersion\",\n",
    "    \"piano music self-expression\",\n",
    "    \"gardening nurturing fulfillment\"\n",
    "]\n",
    "\n",
    "\n",
    "\n",
    "# 计算相似度的函数,返回文本之间的余弦相似度\n",
    "for text in texts:\n",
    "    print(f\"文本: '{text}'\")\n",
    "    scores = []\n",
    "    text_embedding = model.encode(text, convert_to_tensor=True)\n",
    "    for entity_id, entity_text in json_data[\"entities\"].items():\n",
    "        entity_embedding = model.encode(entity_id, convert_to_tensor=True)\n",
    "        similarity_score = util.pytorch_cos_sim(text_embedding, entity_embedding).item()\n",
    "        scores.append((entity_id, similarity_score))\n",
    "\n",
    "    # 根据相似度分数排序并选择最相似的3个实体\n",
    "    scores.sort(key=lambda x: x[1], reverse=True)\n",
    "    top_matches = scores[:3]\n",
    "\n",
    "    # 打印匹配结果及其相似度分数\n",
    "    for match in top_matches:\n",
    "        print(f\"匹配到 '{match[0]}', 相似度分数: {match[1]:.4f}\")\n",
    "    print(\"\\n\" + \"=\"*80 + \"\\n\")\n"
   ]
  }
 ],
 "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
}