fix notebook

This commit is contained in:
Ray 2025-08-26 21:09:15 +08:00
parent 74ec78af36
commit d1cb372b95

View file

@ -94,7 +94,7 @@
"id": "edTfrizMFK4c" "id": "edTfrizMFK4c"
}, },
"source": [ "source": [
"#### 0.1 Install dependencies" "#### 0.1 Install PageIndex"
] ]
}, },
{ {
@ -106,7 +106,7 @@
}, },
"outputs": [], "outputs": [],
"source": [ "source": [
"%pip install -q --upgrade pageindex openai" "%pip install -q --upgrade pageindex"
] ]
}, },
{ {
@ -115,7 +115,7 @@
"id": "WVEWzPKGcG1M" "id": "WVEWzPKGcG1M"
}, },
"source": [ "source": [
"#### 0.2 Setup environment" "#### 0.2 Setup PageIndex"
] ]
}, },
{ {
@ -126,15 +126,40 @@
}, },
"outputs": [], "outputs": [],
"source": [ "source": [
"import json, os, requests\n",
"from pageindex import PageIndexClient\n", "from pageindex import PageIndexClient\n",
"import pageindex.utils as utils\n", "import pageindex.utils as utils\n",
"\n", "\n",
"# Get your PageIndex API key from https://dash.pageindex.ai/api-keys\n", "# Get your PageIndex API key from https://dash.pageindex.ai/api-keys\n",
"PAGEINDEX_API_KEY = \"YOUR_PAGEINDEX_API_KEY\"\n", "PAGEINDEX_API_KEY = \"YOUR_PAGEINDEX_API_KEY\"\n",
"pi_client = PageIndexClient(api_key=PAGEINDEX_API_KEY)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### 0.3 Setup LLM\n",
"\n",
"Choose your preferred LLM for reasoning-based retrieval. In this example, we use OpenAIs GPT-4.1."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import openai\n",
"OPENAI_API_KEY = \"YOUR_OPENAI_API_KEY\"\n", "OPENAI_API_KEY = \"YOUR_OPENAI_API_KEY\"\n",
"\n", "\n",
"pi_client = PageIndexClient(api_key=PAGEINDEX_API_KEY)" "async def call_llm(prompt, model=\"gpt-4.1\", temperature=0):\n",
" client = openai.AsyncOpenAI(api_key=OPENAI_API_KEY)\n",
" response = await client.chat.completions.create(\n",
" model=model,\n",
" messages=[{\"role\": \"user\", \"content\": prompt}],\n",
" temperature=temperature\n",
" )\n",
" return response.choices[0].message.content.strip()"
] ]
}, },
{ {
@ -176,6 +201,8 @@
} }
], ],
"source": [ "source": [
"import os, requests\n",
"\n",
"# You can also use our GitHub repo to generate PageIndex tree\n", "# You can also use our GitHub repo to generate PageIndex tree\n",
"# https://github.com/VectifyAI/PageIndex\n", "# https://github.com/VectifyAI/PageIndex\n",
"\n", "\n",
@ -321,6 +348,8 @@
}, },
"outputs": [], "outputs": [],
"source": [ "source": [
"import json\n",
"\n",
"query = \"What are the conclusions in this document?\"\n", "query = \"What are the conclusions in this document?\"\n",
"\n", "\n",
"tree_without_text = utils.remove_fields(tree.copy(), fields=['text'])\n", "tree_without_text = utils.remove_fields(tree.copy(), fields=['text'])\n",
@ -343,7 +372,7 @@
"Directly return the final JSON structure. Do not output anything else.\n", "Directly return the final JSON structure. Do not output anything else.\n",
"\"\"\"\n", "\"\"\"\n",
"\n", "\n",
"tree_search_result = await utils.call_llm(search_prompt, api_key=OPENAI_API_KEY)" "tree_search_result = await call_llm(search_prompt)"
] ]
}, },
{ {
@ -508,7 +537,7 @@
"\"\"\"\n", "\"\"\"\n",
"\n", "\n",
"print('Generated Answer:\\n')\n", "print('Generated Answer:\\n')\n",
"answer = await utils.call_llm(answer_prompt, api_key=OPENAI_API_KEY)\n", "answer = await call_llm(answer_prompt)\n",
"utils.print_wrapped(answer)" "utils.print_wrapped(answer)"
] ]
}, },