Add GenAI reranker pipeline to llam-rag notebook (#3081)

sbalandi · aleksandr-mokrov · web-flow · commit e992324b1c7e · 2025-09-25T12:50:46.000+02:00
[CVS-170413](https://jira.devtools.intel.com/browse/CVS-170413) --------- Co-authored-by: Aleksandr Mokrov <aleksandr.mokrov@intel.com>
diff --git a/notebooks/llm-rag-langchain/gradio_helper.py b/notebooks/llm-rag-langchain/gradio_helper.py
@@ -40,6 +40,7 @@ def make_demo(
     update_retriever_fn: Callable,
     model_name: str,
     language: str = "English",
+    rerank_device: str | None = None,
 ):
     examples = chinese_examples if (language == "Chinese") else english_examples
 
@@ -226,8 +227,8 @@ def make_demo(
                                 value=2,
                                 step=1,
                                 label="Rerank top n",
-                                info="Number of rerank results",
-                                interactive=True,
+                                info="Number of rerank results(setted on creation step in GenAI pipeline).",
+                                interactive=(rerank_device == "NPU"),
                             )
                         with gr.Row():
                             vector_search_top_k = gr.Slider(
diff --git a/notebooks/llm-rag-langchain/llm-rag-langchain-genai.ipynb b/notebooks/llm-rag-langchain/llm-rag-langchain-genai.ipynb
@@ -147,18 +147,10 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 2,
    "id": "1b2c3f4e",
    "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "LLM config will be updated\n"
-     ]
-    }
-   ],
+   "outputs": [],
    "source": [
     "import os\n",
     "from pathlib import Path\n",
@@ -880,7 +872,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 16,
    "id": "d0bab20b",
    "metadata": {},
    "outputs": [],
@@ -918,14 +910,6 @@
    "id": "e11e73cf",
    "metadata": {},
    "outputs": [
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "D:\\openvino_notebooks\\openvino_env\\Lib\\site-packages\\openvino\\runtime\\__init__.py:10: DeprecationWarning: The `openvino.runtime` module is deprecated and will be removed in the 2026.0 release. Please replace `openvino.runtime` with `openvino`.\n",
-      "  warnings.warn(\n"
-     ]
-    },
     {
      "data": {
       "application/vnd.jupyter.widget-view+json": {
@@ -1198,19 +1182,19 @@
     "\n",
     "[back to top ⬆️](#Table-of-contents:)\n",
     "\n",
-    "Now a local rerank model of OpenVINO can be accelerated on NPU by using the `OpenVINOReranker` class without PyTorch requirements.\n",
+    "Now a local rerank model of OpenVINO can be run via [GenAI `TextRerankPipeline`](https://github.com/openvinotoolkit/openvino.genai/tree/master/samples/python/rag) through `OpenVINOGenAIReranker` and accelerated on NPU by using the `OpenVINOReranker` class without PyTorch requirements.\n",
     "\n",
     "> **Note**: Rerank can be skipped in RAG.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 25,
+   "execution_count": null,
    "id": "b67b39f2-8394-45fb-9b2b-ea63e267a2d3",
    "metadata": {},
    "outputs": [],
    "source": [
-    "from ov_langchain_helper import OpenVINOReranker\n",
+    "from ov_langchain_helper import OpenVINOReranker, OpenVINOGenAIReranker\n",
     "\n",
     "rerank_model_name = rerank_model_id.value\n",
     "rerank_model_kwargs = {\"device_name\": rerank_device.value}\n",
@@ -1235,10 +1219,10 @@
     "        top_n=rerank_top_n,\n",
     "    )\n",
     "else:\n",
-    "    reranker = OpenVINOReranker(\n",
+    "    reranker = OpenVINOGenAIReranker.from_model_path(\n",
     "        model_path=rerank_model_name,\n",
-    "        model_kwargs=rerank_model_kwargs,\n",
-    "        top_n=rerank_top_n,\n",
+    "        device=rerank_device.value,\n",
+    "        config_kwargs={\"top_n\": rerank_top_n},\n",
     "    )"
    ]
   },
@@ -1567,8 +1551,8 @@
     "\n",
     "    if rerank_device.value == \"NPU\":\n",
     "        vector_search_top_k = vector_search_top_k_npu\n",
-    "    if vector_rerank_top_n > vector_search_top_k:\n",
-    "        gr.Warning(\"Search top k must >= Rerank top n\")\n",
+    "        if vector_rerank_top_n > vector_search_top_k:\n",
+    "            gr.Warning(\"Search top k must >= Rerank top n\")\n",
     "\n",
     "    documents = []\n",
     "    for doc in docs:\n",
@@ -1586,7 +1570,8 @@
     "        search_kwargs = {\"k\": vector_search_top_k}\n",
     "    retriever = db.as_retriever(search_kwargs=search_kwargs, search_type=search_method)\n",
     "    if run_rerank:\n",
-    "        reranker.top_n = vector_rerank_top_n\n",
+    "        if rerank_device.value == \"NPU\":\n",
+    "            reranker.top_n = vector_rerank_top_n\n",
     "        retriever = ContextualCompressionRetriever(base_compressor=reranker, base_retriever=retriever)\n",
     "    prompt = PromptTemplate.from_template(rag_prompt_template)\n",
     "    combine_docs_chain = create_stuff_documents_chain(llm, prompt)\n",
@@ -1623,7 +1608,8 @@
     "    retriever = db.as_retriever(search_kwargs=search_kwargs, search_type=search_method)\n",
     "    if run_rerank:\n",
     "        retriever = ContextualCompressionRetriever(base_compressor=reranker, base_retriever=retriever)\n",
-    "        reranker.top_n = vector_rerank_top_n\n",
+    "        if rerank_device.value == \"NPU\":\n",
+    "            reranker.top_n = vector_rerank_top_n\n",
     "    rag_chain = create_retrieval_chain(retriever, combine_docs_chain)\n",
     "\n",
     "    return \"Vector database is Ready\"\n",
@@ -1735,6 +1721,7 @@
     "    update_retriever_fn=update_retriever,\n",
     "    model_name=llm_model_id.value,\n",
     "    language=model_language.value,\n",
+    "    rerank_device=rerank_device.value,\n",
     ")\n",
     "\n",
     "try:\n",
diff --git a/notebooks/llm-rag-langchain/ov_langchain_helper.py b/notebooks/llm-rag-langchain/ov_langchain_helper.py
@@ -736,3 +736,68 @@ def compress_documents(
             )
             final_results.append(doc)
         return final_results
+
+
+class OpenVINOGenAIReranker(BaseDocumentCompressor):
+    """OpenVINO reranking models.
+
+    To use, you should have the ``openvino-genai`` python package installed.
+
+    Example:
+        .. code-block:: python
+
+            from ov_langchain_helper import OpenVINOGenAIReranker
+
+            model_path = "./sentence-transformers/all-mpnet-base-v2"
+            config_kwargs = {'top_n': 3}
+            ov = OpenVINOGenAIReranker.from_model_path(
+                model_path=model_path,
+                device='CPU',
+                config_kwargs=config_kwargs,
+            )
+    """
+
+    ov_pipe: Any = None
+    """OpenVINO pipeline object."""
+    top_n: int = 3
+    """return Top n texts. can not be updated after pipeline was created."""
+
+    model_config = ConfigDict(extra="forbid", protected_namespaces=())
+
+    @classmethod
+    def from_model_path(cls, model_path: str, device: str = "CPU", config_kwargs: Dict[str, Any] = {}) -> OpenVINOGenAIReranker:
+        """Construct the openvino text embedding pipeline from model_path"""
+        try:
+            import openvino_genai
+
+        except ImportError:
+            raise ImportError("Could not import OpenVINO GenAI package. " "Please install it with `pip install openvino-genai`.")
+
+        config = openvino_genai.TextRerankPipeline.Config()
+        if "top_n" in config_kwargs:
+            config.top_n = config_kwargs["top_n"]
+            top_n = config_kwargs["top_n"]
+        if "max_length" in config_kwargs:
+            config.max_length = config_kwargs["max_length"]
+
+        ov_pipe = openvino_genai.TextRerankPipeline(model_path, device, config)
+
+        return cls(ov_pipe=ov_pipe)
+
+    def compress_documents(
+        self,
+        documents: Sequence[Document],
+        query: str,
+        callbacks: Optional[Callbacks] = None,
+    ) -> Sequence[Document]:
+        docs = [doc.page_content for doc in documents]
+        rerank_response = self.ov_pipe.rerank(query, docs)
+        final_results = []
+        for index, score in rerank_response:
+            doc = Document(
+                page_content=documents[index].page_content,
+                metadata={"id": index, "relevance_score": score},
+            )
+            final_results.append(doc)
+
+        return final_results
diff --git a/notebooks/llm-rag-llamaindex/gradio_helper.py b/notebooks/llm-rag-llamaindex/gradio_helper.py
@@ -39,6 +39,7 @@ def make_demo(
     update_retriever_fn: Callable,
     model_name: str,
     language: Literal["English", "Chinese"] = "English",
+    rerank_device: str = "auto",
 ):
     examples = chinese_examples if (language == "Chinese") else english_examples
 
@@ -189,8 +190,8 @@ def make_demo(
                                 value=2,
                                 step=1,
                                 label="Rerank top n",
-                                info="Number of rerank results",
-                                interactive=True,
+                                info="Number of rerank results(setted on creation step in GenAI pipeline).",
+                                interactive=(rerank_device == "NPU"),
                             )
                         with gr.Row():
                             vector_search_top_k = gr.Slider(
diff --git a/notebooks/llm-rag-llamaindex/llm-rag-llamaindex.ipynb b/notebooks/llm-rag-llamaindex/llm-rag-llamaindex.ipynb
@@ -1154,14 +1154,14 @@
     "\n",
     "[back to top ⬆️](#Table-of-contents:)\n",
     "\n",
-    "Now a Hugging Face rerank model can be supported by OpenVINO through [`OpenVINORerank`](https://docs.llamaindex.ai/en/stable/examples/node_postprocessor/openvino_rerank/) class of LlamaIndex.\n",
+    "Now a Hugging Face rerank model can be supported by OpenVINO via [GenAI](https://github.com/openvinotoolkit/openvino.genai) through `OpenVINOGenAIReranking` class. To run model on NPU use [`OpenVINORerank`](https://docs.llamaindex.ai/en/stable/examples/node_postprocessor/openvino_rerank/) class of LlamaIndex.\n",
     "\n",
     "> **Note**: Rerank can be skipped in RAG.\n"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": null,
    "id": "b67b39f2-8394-45fb-9b2b-ea63e267a2d3",
    "metadata": {},
    "outputs": [
@@ -1175,8 +1175,12 @@
    ],
    "source": [
     "from llama_index.postprocessor.openvino_rerank import OpenVINORerank\n",
+    "from ov_llamaindex_helper import OpenVINOGenAIReranking\n",
     "\n",
-    "reranker = OpenVINORerank(model_id_or_path=rerank_model_id.value, device=rerank_device.value, top_n=2)"
+    "if USING_NPU:\n",
+    "    reranker = OpenVINORerank(model_id_or_path=rerank_model_id.value, device=rerank_device.value, top_n=2)\n",
+    "else:\n",
+    "    reranker = OpenVINOGenAIReranking(model_id_or_path=rerank_model_id.value, device=rerank_device.value, top_n=2)"
    ]
   },
   {
@@ -1765,6 +1769,7 @@
     "    update_retriever_fn=update_retriever,\n",
     "    model_name=llm_model_id.value,\n",
     "    language=model_language.value,\n",
+    "    rerank_device=rerank_device.value,\n",
     ")\n",
     "\n",
     "try:\n",
diff --git a/notebooks/llm-rag-llamaindex/ov_llamaindex_helper.py b/notebooks/llm-rag-llamaindex/ov_llamaindex_helper.py
@@ -2,9 +2,21 @@
     DEFAULT_EMBED_BATCH_SIZE,
     BaseEmbedding,
 )
+from llama_index.core.postprocessor.types import BaseNodePostprocessor
 from typing import Any, List, Optional, Dict
 from llama_index.core.bridge.pydantic import Field, PrivateAttr
 from llama_index.core.callbacks import CallbackManager
+from llama_index.core.callbacks import CBEventType, EventPayload
+from llama_index.core.instrumentation import get_dispatcher
+from llama_index.core.instrumentation.events.rerank import (
+    ReRankEndEvent,
+    ReRankStartEvent,
+)
+from llama_index.core.schema import MetadataMode, NodeWithScore, QueryBundle
+from llama_index.core.instrumentation import get_dispatcher
+
+
+dispatcher = get_dispatcher(__name__)
 
 
 class OpenVINOGenAIEmbedding(BaseEmbedding):
@@ -86,3 +98,86 @@ def _get_text_embedding(self, text: str) -> List[float]:
     def _get_text_embeddings(self, texts: List[str]) -> List[List[float]]:
         """Get text embeddings."""
         return self._ov_pipe.embed_documents(texts)
+
+
+class OpenVINOGenAIReranking(BaseNodePostprocessor):
+    model_id_or_path: str = Field(description="Huggingface model id or local path.")
+    top_n: int = Field(description="Number of nodes to return sorted by score.")
+    keep_retrieval_score: bool = Field(
+        default=False,
+        description="Whether to keep the retrieval score in metadata.",
+    )
+
+    _ov_pipe: Any = PrivateAttr()
+
+    def __init__(
+        self,
+        model_id_or_path: str,
+        max_length: Optional[int] = None,
+        top_n: Optional[int] = 3,
+        device: Optional[str] = "auto",
+        model_kwargs: Dict[str, Any] = {},
+        keep_retrieval_score: Optional[bool] = False,
+    ):
+        try:
+            import openvino_genai
+        except ImportError:
+            raise ImportError("Could not import OpenVINO GenAI package. " "Please install it with `pip install openvino-genai`.")
+
+        super().__init__(top_n=top_n, max_length=max_length, model_id_or_path=model_id_or_path, device=device, keep_retrieval_score=keep_retrieval_score)
+
+        config = openvino_genai.TextRerankPipeline.Config()
+        config.top_n = top_n
+        if max_length:
+            config.max_length = max_length
+
+        ov_pipe = openvino_genai.TextRerankPipeline(model_id_or_path, device, config, **model_kwargs)
+
+        self._ov_pipe = ov_pipe
+
+    @classmethod
+    def class_name(cls) -> str:
+        return "OpenVINOGenAIReranking"
+
+    def _postprocess_nodes(
+        self,
+        nodes: List[NodeWithScore],
+        query_bundle: Optional[QueryBundle] = None,
+    ) -> List[NodeWithScore]:
+        dispatcher.event(
+            ReRankStartEvent(
+                query=query_bundle,
+                nodes=nodes,
+                top_n=self.top_n,
+                model_name=self.model_id_or_path,
+            )
+        )
+
+        if query_bundle is None:
+            raise ValueError("Missing query bundle in extra info.")
+        if len(nodes) == 0:
+            return []
+
+        nodes_text_list = [str(node.node.get_content(metadata_mode=MetadataMode.EMBED)) for node in nodes]
+
+        with self.callback_manager.event(
+            CBEventType.RERANKING,
+            payload={
+                EventPayload.NODES: nodes,
+                EventPayload.MODEL_NAME: self.model_id_or_path,
+                EventPayload.QUERY_STR: query_bundle.query_str,
+                EventPayload.TOP_K: self.top_n,
+            },
+        ) as event:
+            outputs = self._ov_pipe.rerank(query_bundle.query_str, nodes_text_list)
+
+            for node, score in zip(nodes, outputs):
+                if self.keep_retrieval_score:
+                    # keep the retrieval score in metadata
+                    node.node.metadata["retrieval_score"] = node.score
+                node.score = score
+
+            event.on_end(payload={EventPayload.NODES: nodes})
+
+        dispatcher.event(ReRankEndEvent(nodes=nodes))
+        return nodes
