File size: 1,024 Bytes
407c075
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores.faiss import FAISS


class DocumentStore:

  def __getitem__(self, query):
    raise NotImplementedError


class FaissDocumentStore(DocumentStore):
  filename = "documents/stored.pkl"

  def __init__(self, store):
    self.store = store

  @classmethod
  def from_texts(cls, texts, metadatas):
    embeddings = OpenAIEmbeddings()  # uses ada-002 by default
    docsearch = FAISS.from_texts(texts, embeddings, metadatas=metadatas)

    return cls(docsearch)

  @classmethod
  def from_pickle(cls, filename=None):
    import pickle

    if filename is None:
      filename = cls.filename

    with open(filename, "rb") as f:
      store = pickle.load(f)

    return cls(store)

  def to_pickle(self, filename=None):
    import pickle

    if filename is None:
      filename = self.filename

    with open(filename, "wb") as f:
      pickle.dump(self.store, f)

  def __getitem__(self, query):
    return self.store.similarity_search(query)