Skip to content

Commit

Permalink
feat: add Jina Reranking integration (#11291)
Browse files Browse the repository at this point in the history 
* feat: add jina reranking

Signed-off-by: Joan Martinez <[email protected]>

* feat: add notebook

* fix: fix linting

* fix: change printing of nodes

* fix: change printing of nodes

* fix: fix output from new model

* fix: fix notebook

* cr

---------

Signed-off-by: Joan Martinez <[email protected]>
Co-authored-by: Haotian Zhang <[email protected]>
  • Loading branch information
@JoanFM @hatianzhang
JoanFM and hatianzhang authored Feb 28, 2024
1 parent 841dca5 commit 8623877
Show file tree
Hide file tree
Showing 12 changed files with 664 additions and 0 deletions.
325 changes: 325 additions & 0 deletions docs/examples/node_postprocessor/JinaRerank.ipynb
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,325 @@
{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/node_postprocessor/CohereRerank.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Jina Rerank"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"!pip install llama-index-postprocessor-jinaai-rerank\n",
"!pip install llama-index-embeddings-jinaai\n",
"!pip install llama-index"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from llama_index.core import (\n",
" VectorStoreIndex,\n",
" SimpleDirectoryReader,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from llama_index.embeddings.jinaai import JinaEmbedding\n",
"\n",
"api_key = os.environ[\"JINA_API_KEY\"]\n",
"jina_embeddings = JinaEmbedding(api_key=api_key)\n",
"\n",
"# load documents\n",
"import requests\n",
"\n",
"url = \"https://niketeam-asset-download.nike.net/catalogs/2024/2024_Nike%20Kids_02_09_24.pdf?cb=09302022\"\n",
"response = requests.get(url)\n",
"with open(\"Nike_Catalog.pdf\", \"wb\") as f:\n",
" f.write(response.content)\n",
"reader = SimpleDirectoryReader(input_files=[\"Nike_Catalog.pdf\"])\n",
"documents = reader.load_data()\n",
"\n",
"# build index\n",
"index = VectorStoreIndex.from_documents(\n",
" documents=documents, embed_model=jina_embeddings\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Retrieve top 10 most relevant nodes, without using a reranker"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"query_engine = index.as_query_engine(similarity_top_k=10)\n",
"response = query_engine.query(\n",
" \"What is the best jersey by Nike in terms of fabric?\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"56\n",
"Sustainable MaterialsNIKE KIDS SOCCER – GOALKEEPER\n",
"KIDS NIKE DRY LS US PARK IV GK JERSEY \n",
"CJ6073 $42.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 04/01/20\n",
"END DATE: 12/31/25\n",
"Goal keepers jersey with graphic print on sleeves and across upper back panel, mesh back for breathability, \n",
"slim fit with soft hand feel, shoulder seams rolled forward for better graphic visibility, straight seam across \n",
"back, mesh back for breathability – gameday graphic print inspired by retro campos gk design . \n",
"Body width: 16.3\", Body length: 22\" (size medium).\n",
"010 Black/White/(White) 012 Wolf Grey/White/(Black) 702 Volt/White/(Black)\n",
"KIDS NIKE DRY PARK III SHORT \n",
"BV6866 $20.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 04/01/20\n",
"END DATE: 12/31/25\n",
"Dri-FIT angled side seam short (slim fit) with soft hand feel updated fit for better mobility/comfort . \n",
"Hip width: 16.9\", Inseam length: 7\" (size medium).\n",
"010 Black/White/(White) 012 Wolf Grey/Black/(Black) 702 Volt/(Black)\n",
"NIKE ACADEMY OTC SOCK (UNISEX) \n",
"SX5728 $12.00\n",
"Sold in prepacks of 6.\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 93% nylon/6% polyester/1% spandex.\n",
"OFFER DATE: 01/01/17\n",
"END DATE: 12/31/23\n",
"Game day sock with fold-over cuff, articulated foot specific footbed for superior fit and contrast Swoosh \n",
"design trademark at ankle. Sold in prepacks of 6.\n",
"010 Black/(White) 018 Wolf Grey/(Black) 702 Volt/(Black)\n",
"Sustainable Materials 0.8641328028479249\n",
"\n",
"\n",
"NIKE KIDS SOCCER – STOCK42\n",
"Sustainable Materials\n",
"KIDS NIKE DRI-FIT US SS \n",
"CHALLENGE IV JERSEY\n",
"DH8368 $42.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 01/01/22\n",
"END DATE: 12/31/23\n",
"The Nike Dri-FIT Challenge IV Jersey brings subtle style and modern performance to the field. Sweat-\n",
"wicking fabric helps keep you dry and comfortable from the first whistle to the last minute.\n",
"010 Black/Black/White/(White) 012 Wolf Grey/Wolf Grey/Black/(Black)\n",
"100 White/White/White/(Black) 341 Gorge Green/Gorge Green/White/(White)\n",
"419 College Navy/College Navy/White/(White) 448 Valor Blue/Valor Blue/White/(White)\n",
"480 Game Royal/Game Royal/White/(White) 657 University Red/University Red/White/(White)\n",
"692 Team Maroon/Team Maroon/White/(White) 702 Volt/Volt/Black/(Black)\n",
"891 Team Orange/Team Orange/Black/(Black)\n",
"NEW\n",
"KIDS NIKE DRI-FIT CHALLENGE V JERSEY \n",
"SS US\n",
"FD7427 $47.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 01/01/24\n",
"END DATE: 12/31/25\n",
"The Nike Dri-FIT Challenge Jersey V is designed to keep your players cool and comfortable through 90 \n",
"minutes and beyond. Mesh on the back and side panels offer breathability where athletes need it most. \n",
"Body and sleeves are a Nike Dri-FIT knit fabric that moves sweat away to help keep players dry. This top \n",
"is made with 100% recycled material. Side panel construction uses a more efficient pattern to help reduce \n",
"material waste. Slim fit for a tailored look and feel.\n",
"010 Black/White/(White) 012 Wolf Grey/Black/(Black) 100 White/Black/(Black)\n",
"341 Gorge Green/White/(White) 419 College Navy/White/(White) 448 Valor Blue/White/(White)\n",
"480 Game Royal/White/(White) 657 University Red/White/(White) 692 Team Maroon/White/(White)\n",
"702 Volt/Black/(Black) 891 Team Orange/Black/(Black)\n",
"BACK VIEW 0.863721033128725\n"
]
}
],
"source": [
"print(response.source_nodes[0].text, response.source_nodes[0].score)\n",
"print(\"\\n\")\n",
"print(response.source_nodes[1].text, response.source_nodes[1].score)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Retrieve top 10 most relevant nodes, but then rerank using Jina Reranker"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"By employing a reranker model, the prompt can be given more relevant context. This will lead to a more accurate response by the LLM."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"from llama_index.postprocessor.jinaai_rerank import JinaRerank\n",
"\n",
"jina_rerank = JinaRerank(api_key=api_key, top_n=2)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"query_engine = index.as_query_engine(\n",
" similarity_top_k=10, node_postprocessors=[jina_rerank]\n",
")\n",
"response = query_engine.query(\n",
" \"What is the best jersey by Nike in terms of fabric?\",\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"NIKE KIDS SOCCER – STOCK41Sustainable Materials\n",
"Sustainable Materials\n",
"KIDS DRI-FIT ADV VAPOR IV JERSEY US SS\n",
"DR0837 $77.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 01/01/23\n",
"END DATE: 12/31/24\n",
"Step on to the field ready for fast-paced play in the Nike Dri-FIT ADV Vapor Jersey. Engineered for \n",
"optimal breathability, its moisture-wicking design helps keep you dry and cool under match-day pressure. \n",
"Lightweight fabric in a relaxed, easy fit combats cling so you can focus on being the first to the ball. Lower \n",
"insets line up perfectly with design details on the Nike Dri-FIT ADV Vapor IV Shorts to create an on-field \n",
"look worthy of pro-level play. \n",
"010 Black/Black/Black/(White) 100 White/White/White/(Black)\n",
"419 College Navy/College Navy/Game Royal/(White) 480 Game Royal/Game Royal/College Navy/(White)\n",
"657 University Red/University Red/Bright Crimson/(White)\n",
"BACK VIEW\n",
"GRAPHIC KNIT DETAIL\n",
"KIDS NIKE DRI-FIT US SS STRIKE III JERSEY\n",
"DR0913 $50.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 01/01/23\n",
"END DATE: 12/31/24\n",
"Take the field in match-ready style in the lightweight Nike Strike Jersey. A relaxed, easy fit ensures that \n",
"nothing comes between you and the ball, and sweat-wicking fabric works with breathable mesh to help \n",
"keep you cool and composed during fast-paced play. Ribbed insets stretch with you to let you move without \n",
"restrictions. Embroidered Swoosh design trademark. \n",
"010 Black/Black/Black/(White) 011 Black/Volt/Volt/(White)\n",
"012 Wolf Grey/Black/Black/(White) 100 White/White/White/(Black)\n",
"419 College Navy/College Navy/Game Royal/(White) 448 Valor Blue/College Navy/College Navy/(White)\n",
"480 Game Royal/College Navy/College Navy/(White) 657 University Red/Bright Crimson/Bright Crimson/(White)\n",
"GRAPHIC KNIT DETAIL 0.3603765070438385\n",
"\n",
"\n",
"NIKE KIDS SOCCER – STOCK45\n",
"Sustainable MaterialsKIDS NIKE DRI-FIT US LS TIEMPO\n",
"PREMIER II JERSEY\n",
"DH8407 $32.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 01/01/22\n",
"END DATE: 12/31/26\n",
"The Nike Dri-FIT Tiempo Premier II Jersey brings you the cool performance of sweat-wicking fabric and a \n",
"mesh back panel kick in when the game heats up.\n",
"010 Black/White/(White) 100 White/White/(Black) 419 College Navy/White/(White)\n",
"480 Game Royal/White/(White) 657 University Red/White/(White)\n",
"KIDS NIKE DRI-FIT US SS TIEMPO\n",
"PREMIER II JERSEY\n",
"DH8390 $27.00\n",
"SIZES: XS, S, M, L, XL\n",
"FABRIC: 100% polyester.\n",
"OFFER DATE: 01/01/22\n",
"END DATE: 12/31/26\n",
"The Nike Dri-FIT Tiempo Premier II Jersey brings you the cool performance of sweat-wicking fabric and a \n",
"mesh back panel kick in when the game heats up.\n",
"010 Black/White/(White) 012 Wolf Grey/Black/(Black) 100 White/White/(Black)\n",
"341 Gorge Green/White/(White) 419 College Navy/White/(White) 448 Valor Blue/White/(White)\n",
"480 Game Royal/White/(White) 547 Court Purple/White/(White) 616 Vivid Pink/Black/(Black)\n",
"657 University Red/White/(White) 692 Team Maroon/White/(White) 702 Volt/Black/(Black)\n",
"891 Team Orange/Black/(Black)\n",
"Sustainable Materials 0.35767972469329834\n"
]
}
],
"source": [
"print(response.source_nodes[0].text, response.source_nodes[0].score)\n",
"print(\"\\n\")\n",
"print(response.source_nodes[1].text, response.source_nodes[1].score)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"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"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
Loading

0 comments on commit 8623877

Please sign in to comment.