add openvino rerank import test

fix format fix format fix format fix format fix format fix format fix format
langchain-ai · Mar 30, 2024 · 809fa81 · 809fa81
1 parent e2b40e2
commit 809fa81
Show file tree

Hide file tree

Showing 5 changed files with 20 additions and 11 deletions.
diff --git a/docs/docs/integrations/retrievers/openvino_rerank.ipynb b/docs/docs/integrations/retrievers/openvino_rerank.ipynb
@@ -82,6 +82,7 @@
    "source": [
     "# Helper function for printing docs\n",
     "\n",
+    "\n",
     "def pretty_print_docs(docs):\n",
     "    print(\n",
     "        f\"\\n{'-' * 100}\\n\".join(\n",
@@ -365,24 +366,23 @@
     }
    ],
    "source": [
+    "from langchain.embeddings import OpenVINOEmbeddings\n",
     "from langchain_community.document_loaders import TextLoader\n",
     "from langchain_community.vectorstores import FAISS\n",
-    "from langchain.embeddings import OpenVINOEmbeddings\n",
     "from langchain_text_splitters import RecursiveCharacterTextSplitter\n",
     "\n",
     "documents = TextLoader(\n",
     "    \"../../modules/state_of_the_union.txt\",\n",
     ").load()\n",
-    "text_splitter = RecursiveCharacterTextSplitter(\n",
-    "    chunk_size=500, chunk_overlap=100)\n",
+    "text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)\n",
     "texts = text_splitter.split_documents(documents)\n",
     "for idx, text in enumerate(texts):\n",
     "    text.metadata[\"id\"] = idx\n",
     "\n",
     "embedding = OpenVINOEmbeddings(\n",
-    "    model_name_or_path=\"sentence-transformers/all-mpnet-base-v2\")\n",
-    "retriever = FAISS.from_documents(\n",
-    "    texts, embedding).as_retriever(search_kwargs={\"k\": 20})\n",
+    "    model_name_or_path=\"sentence-transformers/all-mpnet-base-v2\"\n",
+    ")\n",
+    "retriever = FAISS.from_documents(texts, embedding).as_retriever(search_kwargs={\"k\": 20})\n",
     "\n",
     "query = \"What did the president say about Ketanji Brown Jackson\"\n",
     "docs = retriever.get_relevant_documents(query)\n",
@@ -550,11 +550,15 @@
    ],
    "source": [
     "from pathlib import Path\n",
+    "\n",
     "ov_model_dir = \"bge-reranker-large-ov\"\n",
     "if not Path(ov_model_dir).exists():\n",
     "    from optimum.intel.openvino import OVModelForSequenceClassification\n",
     "    from transformers import AutoTokenizer\n",
-    "    ov_model = OVModelForSequenceClassification.from_pretrained(model_name, compile=False, export=True)\n",
+    "\n",
+    "    ov_model = OVModelForSequenceClassification.from_pretrained(\n",
+    "        model_name, compile=False, export=True\n",
+    "    )\n",
     "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
     "    ov_model.half()\n",
     "    ov_model.save_pretrained(ov_model_dir)\n",
@@ -575,7 +579,7 @@
     }
    ],
    "source": [
-    "compressor = OpenVINOReranker(model_name_or_path=ov_model_dir)\n"
+    "compressor = OpenVINOReranker(model_name_or_path=ov_model_dir)"
    ]
   },
   {

diff --git a/docs/docs/integrations/text_embedding/openvino.ipynb b/docs/docs/integrations/text_embedding/openvino.ipynb
@@ -156,11 +156,15 @@
    "outputs": [],
    "source": [
     "from pathlib import Path\n",
+    "\n",
     "ov_model_dir = \"all-mpnet-base-v2-ov\"\n",
     "if not Path(ov_model_dir).exists():\n",
     "    from optimum.intel.openvino import OVModelForFeatureExtraction\n",
     "    from transformers import AutoTokenizer\n",
-    "    ov_model = OVModelForFeatureExtraction.from_pretrained(model_name, compile=False, export=True)\n",
+    "\n",
+    "    ov_model = OVModelForFeatureExtraction.from_pretrained(\n",
+    "        model_name, compile=False, export=True\n",
+    "    )\n",
     "    tokenizer = AutoTokenizer.from_pretrained(model_name)\n",
     "    ov_model.half()\n",
     "    ov_model.save_pretrained(ov_model_dir)\n",

diff --git a/libs/community/langchain_community/retrievers/openvino_rerank.py b/libs/community/langchain_community/retrievers/openvino_rerank.py
@@ -107,7 +107,7 @@ def require_model_export(
 
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path)
 
-    def rerank(self, request: Any):
+    def rerank(self, request: Any) -> Any:
         query = request.query
         passages = request.passages
 

diff --git a/libs/community/tests/unit_tests/retrievers/test_imports.py b/libs/community/tests/unit_tests/retrievers/test_imports.py
@@ -23,6 +23,7 @@
     "MetalRetriever",
     "MilvusRetriever",
     "OutlineRetriever",
+    "OpenVINOReranker",
     "PineconeHybridSearchRetriever",
     "PubMedRetriever",
     "QdrantSparseVectorRetriever",

diff --git a/libs/langchain/langchain/retrievers/document_compressors/openvino_rerank.py b/libs/langchain/langchain/retrievers/document_compressors/openvino_rerank.py
@@ -2,4 +2,4 @@
 
 __all__ = [
     "OpenVINOReranker",
-]
+]