一、簡介
Langchain-ChatGLM 相信大家都不陌生,近幾周計劃出一個原始碼解讀,先解鎖langchain的一些基礎用法。
文件問答過程大概分為以下5部分,在Langchain中都有體現。
- 上傳解析文件
- 文件向量化、儲存
- 文件召回
- query向量化
- 文件問答
今天主要講langchain在文件embedding以及構建faiss過程時是怎麼實現的。
二、原始碼入口
langchain中對於文件embedding以及構建faiss過程有2個分支,
1.當第一次進行載入檔案時如何生成faiss.index
2.當存在faiss.index時
下面也分別從這2個方面進行原始碼解讀
if len(docs) > 0:
logger.info("檔案載入完畢,正在生成向量庫")
if vs_path and os.path.isdir(vs_path) and "index.faiss" in os.listdir(vs_path):
vector_store = load_vector_store(vs_path, self.embeddings)
vector_store.add_documents(docs)
torch_gc()
else:
if not vs_path:
vs_path = os.path.join(KB_ROOT_PATH,f"""{"".join(lazy_pinyin(os.path.splitext(file)[0]))}_FAISS_{datetime.datetime.now().strftime("%Y%m%d_%H%M%S")}""","vector_store")
vector_store = MyFAISS.from_documents(docs, self.embeddings) # docs 為Document列表
torch_gc()
vector_store.save_local(vs_path)
三、不存在faiss.index
MyFAISS.from_documents()過載了父類VectorStore的from_documents(),這裡的self.embeddings其實是一個embedding物件
vector_store = MyFAISS.from_documents(docs, self.embeddings) # docs 為Document列表
self.embeddings = HuggingFaceEmbeddings(model_name=embedding_model_dict[embedding_model],
model_kwargs={'device': embedding_device})
@classmethod
def from_documents(
cls: Type[VST],
documents: List[Document],
embedding: Embeddings,
**kwargs: Any,) -> VST:
"""Return VectorStore initialized from documents and embeddings."""
texts = [d.page_content for d in documents]
metadatas = [d.metadata for d in documents]
return cls.from_texts(texts, embedding, metadatas=metadatas, **kwargs)
from_texts()主要做了3件事情
1.文件embedding化。
2. 建立記憶體中的文件儲存
3.初始化FAISS資料庫
最後返回的cls例項指class FAISS(VectorStore):
在這裡,要注意2個變數,
embeddings: List[List[float]], 真正的embedding向量,2維列表
embedding: Embeddings, 一個huggingface類物件
@classmethod
def from_texts(
cls,
texts: List[str],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
**kwargs: Any,) -> FAISS:
"""Construct FAISS wrapper from raw documents. This is a user friendly interface that: 1. Embeds documents. 2. Creates an in memory docstore 3. Initializes the FAISS database This is intended to be a quick way to get started. Example: .. code-block:: python from langchain import FAISS from langchain.embeddings import OpenAIEmbeddings embeddings = OpenAIEmbeddings() faiss = FAISS.from_texts(texts, embeddings) """
embeddings = embedding.embed_documents(texts)
return cls.__from(
texts,
embeddings,
embedding,
metadatas,
**kwargs, )
現在先看embeddings = embedding.embed_documents(texts)
我們在chains/modules/embeddings.py找到,可以看到embedding依賴client.encode這個函式,所以說如果想要自定義embedding模型,如果是huggingface的,那就比較簡單,定義好模型名稱和模型路徑就可以了。這裡返回一個向量List[List[float]]
def embed_documents(self, texts: List[str]) -> List[List[float]]:
"""Compute doc embeddings using a HuggingFace transformer model. Args: texts: The list of texts to embed. Returns: List of embeddings, one for each text. """
texts = list(map(lambda x: x.replace("\n", " "), texts))
embeddings = self.client.encode(texts, normalize_embeddings=True)
return embeddings.tolist()
self.client = sentence_transformers.SentenceTransformer(
self.model_name, cache_folder=self.cache_folder, **self.model_kwargs
)
接下來看cls.__from(),可以看出faiss構建索引的核心在這裡面
@classmethod
def __from(
cls,
texts: List[str],
embeddings: List[List[float]],
embedding: Embeddings,
metadatas: Optional[List[dict]] = None,
normalize_L2: bool = False,
**kwargs: Any,) -> FAISS:
faiss = dependable_faiss_import()
index = faiss.IndexFlatL2(len(embeddings[0]))
vector = np.array(embeddings, dtype=np.float32)
if normalize_L2:
faiss.normalize_L2(vector)
index.add(vector)
documents = []
for i, text in enumerate(texts):
metadata = metadatas[i] if metadatas else {}
documents.append(Document(page_content=text, metadata=metadata))
index_to_id = {i: str(uuid.uuid4()) for i in range(len(documents))}
docstore = InMemoryDocstore(
{index_to_id[i]: doc for i, doc in enumerate(documents)}
)
return cls(
embedding.embed_query,
index,
docstore,
index_to_id,
normalize_L2=normalize_L2,
**kwargs, )
從中我們看到幾個faiss構建索引的要素
1.faiss.IndexFlatL2,L2衡量距離
2.index_to_id ,給每個chunk一個獨一無二的編碼id,{hashcode:chunk1, ... ...}
3.docstore,其實是一個InMemoryDocstore類
class InMemoryDocstore(Docstore, AddableMixin):
"""Simple in memory docstore in the form of a dict."""
def __init__(self, _dict: Dict[str, Document]):
"""Initialize with dict."""
self._dict = _dict
def add(self, texts: Dict[str, Document]) -> None:
"""Add texts to in memory dictionary."""
overlapping = set(texts).intersection(self._dict)
if overlapping:
raise ValueError(f"Tried to add ids that already exist: {overlapping}")
self._dict = dict(self._dict, **texts)
def search(self, search: str) -> Union[str, Document]:
"""Search via direct lookup."""
if search not in self._dict:
return f"ID {search} not found."
else:
return self._dict[search]
4.embed_query(),這是HuggingFace的一個encoding的方法,這裡真正實現了把文字進行向量化的過程
def embed_query(self, text: str) -> List[float]:
"""
Compute query embeddings using a HuggingFace transformer model.
Args:
text: The text to embed.
Returns: Embeddings for the text. """
text = text.replace("\n", " ")
embedding = self.client.encode(text, **self.encode_kwargs)
return embedding.tolist()
最後可以看到vector_store其實就是一個包含文件資訊的FAISS物件,其中向量化的過程已經在流程中生成了檔案
vector_store = MyFAISS.from_documents(docs, self.embeddings) # docs 為Document列表
class FAISS(VectorStore):
"""Wrapper around FAISS vector database. To use, you should have the ``faiss`` python package installed. Example: .. code-block:: python from langchain import FAISS faiss = FAISS(embedding_function, index, docstore, index_to_docstore_id) """
def __init__(
self,
embedding_function: Callable,
index: Any,
docstore: Docstore,
index_to_docstore_id: Dict[int, str],
relevance_score_fn: Optional[
Callable[[float], float]
] = _default_relevance_score_fn,
normalize_L2: bool = False, ):
三、存在faiss.index
vector_store = load_vector_store(vs_path, self.embeddings)
這裡做了lru_cache快取機制, MyFAISS呼叫靜態方法load_local
@lru_cache(CACHED_VS_NUM)
def load_vector_store(vs_path, embeddings):
return MyFAISS.load_local(vs_path, embeddings)
可以看到最後返回的是一個vector_store的FAISS(VectorStore)類
@classmethoddef load_local(
cls, folder_path: str, embeddings: Embeddings, index_name: str = "index") -> FAISS:
"""Load FAISS index, docstore, and index_to_docstore_id to disk. Args: folder_path: folder path to load index, docstore, and index_to_docstore_id from. embeddings: Embeddings to use when generating queries index_name: for saving with a specific index file name """
path = Path(folder_path)
# load index separately since it is not picklable faiss = dependable_faiss_import()
index = faiss.read_index(
str(path / "{index_name}.faiss".format(index_name=index_name))
)
# load docstore and index_to_docstore_id with open(path / "{index_name}.pkl".format(index_name=index_name), "rb") as f:
docstore, index_to_docstore_id = pickle.load(f)
return cls(embeddings.embed_query, index, docstore, index_to_docstore_id)
進入主題,vector_store.add_documents(docs)巢狀了2個函式,依次如下
def add_documents(self, documents: List[Document], **kwargs: Any) -> List[str]:
"""Run more documents through the embeddings and add to the vectorstore. Args: documents (List[Document]: Documents to add to the vectorstore. Returns: List[str]: List of IDs of the added texts. """ # TODO: Handle the case where the user doesn't provide ids on the Collection texts = [doc.page_content for doc in documents]
metadatas = [doc.metadata for doc in documents]
return self.add_texts(texts, metadatas, **kwargs)
def add_texts(
self,
texts: Iterable[str],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,) -> List[str]:
"""Run more texts through the embeddings and add to the vectorstore. Args: texts: Iterable of strings to add to the vectorstore. metadatas: Optional list of metadatas associated with the texts. Returns: List of ids from adding the texts into the vectorstore. """
if not isinstance(self.docstore, AddableMixin):
raise ValueError(
"If trying to add texts, the underlying docstore should support " f"adding items, which {self.docstore} does not" )
# Embed and create the documents. embeddings = [self.embedding_function(text) for text in texts]
return self.__add(texts, embeddings, metadatas, **kwargs)
def __add(
self,
texts: Iterable[str],
embeddings: Iterable[List[float]],
metadatas: Optional[List[dict]] = None,
**kwargs: Any,) -> List[str]:
if not isinstance(self.docstore, AddableMixin):
raise ValueError(
"If trying to add texts, the underlying docstore should support " f"adding items, which {self.docstore} does not" )
documents = []
for i, text in enumerate(texts):
metadata = metadatas[i] if metadatas else {}
documents.append(Document(page_content=text, metadata=metadata))
# Add to the index, the index_to_id mapping, and the docstore. starting_len = len(self.index_to_docstore_id)
faiss = dependable_faiss_import()
vector = np.array(embeddings, dtype=np.float32)
if self._normalize_L2:
faiss.normalize_L2(vector)
self.index.add(vector)
# Get list of index, id, and docs. full_info = [
(starting_len + i, str(uuid.uuid4()), doc)
for i, doc in enumerate(documents)
]
# Add information to docstore and index. self.docstore.add({_id: doc for _, _id, doc in full_info})
index_to_id = {index: _id for index, _id, _ in full_info}
self.index_to_docstore_id.update(index_to_id)
return [_id for _, _id, _ in full_info]
其中,self.index做的是向量增量操作;full_info,self.docstore,self.index_to_docstore_id做的都是資料增量操作,add_documents()返回如下,可以看出是一個文件編碼list
full_info = [
(starting_len + i, str(uuid.uuid4()), doc)
for i, doc in enumerate(documents)
]
return [_id for _, _id, _ in full_info]
四、檔案儲存
檔案儲存,存了index、docstore、index_to_docstore_id,其中
{index_name}.faiss:向量儲存的faiss索引
{index_name}.pkl:存取的docstore物件以及index_to_docstore_id字典
def save_local(self, folder_path: str, index_name: str = "index") -> None:
"""Save FAISS index, docstore, and index_to_docstore_id to disk. Args: folder_path: folder path to save index, docstore, and index_to_docstore_id to. index_name: for saving with a specific index file name """
path = Path(folder_path)
path.mkdir(exist_ok=True, parents=True)
# save index separately since it is not picklable faiss = dependable_faiss_import()
faiss.write_index(
self.index, str(path / "{index_name}.faiss".format(index_name=index_name))
)
# save docstore and index_to_docstore_id
with open(path / "{index_name}.pkl".format(index_name=index_name), "wb") as f:
pickle.dump((self.docstore, self.index_to_docstore_id), f)
五、雜談
文件embedding以及構建faiss過程在實現中其實很繞,需要用心去讀原始碼,還有關注函式define中的變數型別,理解才會事半功倍!
胖友,請不要忘了一鍵三連點贊哦!
轉載請註明出處:QA Weekly