From df0c33a005c6f33163da633480d51cbae33569f4 Mon Sep 17 00:00:00 2001
From: Davis Chase <130488702+dev2049@users.noreply.github.com>
Date: Wed, 17 May 2023 19:18:57 -0700
Subject: [PATCH] Faiss no avx2 (#4895)

Co-authored-by: Ali Mirlou <alimirlou@gmail.com>
---
 .../indexes/vectorstores/examples/faiss.ipynb |  5 +++-
 langchain/vectorstores/faiss.py               | 23 +++++++++++++++----
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/docs/modules/indexes/vectorstores/examples/faiss.ipynb b/docs/modules/indexes/vectorstores/examples/faiss.ipynb
index 6482efc1..78000c2c 100644
--- a/docs/modules/indexes/vectorstores/examples/faiss.ipynb
+++ b/docs/modules/indexes/vectorstores/examples/faiss.ipynb
@@ -56,7 +56,10 @@
     "import os\n",
     "import getpass\n",
     "\n",
-    "os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')"
+    "os.environ['OPENAI_API_KEY'] = getpass.getpass('OpenAI API Key:')\n",
+    "\n",
+    "# Uncomment the following line if you need to initialize FAISS with no AVX2 optimization\n",
+    "# os.environ['FAISS_NO_AVX2'] = '1'"
    ]
   },
   {
diff --git a/langchain/vectorstores/faiss.py b/langchain/vectorstores/faiss.py
index a606174a..2085f95a 100644
--- a/langchain/vectorstores/faiss.py
+++ b/langchain/vectorstores/faiss.py
@@ -2,6 +2,7 @@
 from __future__ import annotations
 
 import math
+import os
 import pickle
 import uuid
 from pathlib import Path
@@ -17,10 +18,24 @@ from langchain.vectorstores.base import VectorStore
 from langchain.vectorstores.utils import maximal_marginal_relevance
 
 
-def dependable_faiss_import() -> Any:
-    """Import faiss if available, otherwise raise error."""
+def dependable_faiss_import(no_avx2: Optional[bool] = None) -> Any:
+    """
+    Import faiss if available, otherwise raise error.
+    If FAISS_NO_AVX2 environment variable is set, it will be considered
+    to load FAISS with no AVX2 optimization.
+
+    Args:
+        no_avx2: Load FAISS strictly with no AVX2 optimization
+            so that the vectorstore is portable and compatible with other devices.
+    """
+    if no_avx2 is None and "FAISS_NO_AVX2" in os.environ:
+        no_avx2 = bool(os.getenv("FAISS_NO_AVX2"))
+
     try:
-        import faiss
+        if no_avx2:
+            from faiss import swigfaiss as faiss
+        else:
+            import faiss
     except ImportError:
         raise ValueError(
             "Could not import faiss python package. "
@@ -161,7 +176,7 @@ class FAISS(VectorStore):
         """Return docs most similar to query.
 
         Args:
-            query: Text to look up documents similar to.
+            embedding: Embedding vector to look up documents similar to.
             k: Number of Documents to return. Defaults to 4.
 
         Returns: