similarity.py

import os, sys
import re,time
from os import path

sys.path.append("../") 
import copy
from typing import List,OrderedDict,Any,Optional,Tuple,Dict
from vector.pgsqldocstore import InMemorySecondaryDocstore
from langchain.vectorstores.faiss import FAISS,dependable_faiss_import
from langchain.schema import Document
from vector.pgsqldocstore import PgSqlDocstore   
from langchain.embeddings.huggingface import (
    HuggingFaceEmbeddings,
)
import math
import faiss
from langchain.vectorstores.utils import DistanceStrategy
from langchain.vectorstores.base import VectorStoreRetriever
from langchain.callbacks.manager import (
    AsyncCallbackManagerForRetrieverRun,
    CallbackManagerForRetrieverRun,
)
from loader import load
from langchain.embeddings.base import Embeddings
from vector.VectorCallback import DocumentCallback,DefaultDocumentCallback

def singleton(cls):
    instances = {}
    def get_instance(*args, **kwargs):
        if cls not in instances:
            instances[cls] = cls(*args, **kwargs)
        return instances[cls]
    return get_instance

@singleton
class EmbeddingFactory:
    def __init__(self, path:str):
        self.path = path
        self.embedding = HuggingFaceEmbeddings(model_name=path)

    def get_embedding(self):
        return self.embedding

def GetEmbding(path:str) -> Embeddings:
    # return HuggingFaceEmbeddings(model_name=path)
    return EmbeddingFactory(path).get_embedding()
    
import operator
from langchain.vectorstores.utils import DistanceStrategy
import numpy as np
class RE_FAISS(FAISS):
    #去重，并保留metadate
    def _tuple_deduplication(self, tuple_input:List[Tuple[Document, float]]) -> List[Tuple[Document, float]]:
        deduplicated_dict = OrderedDict()
        for doc,scores in tuple_input:
            page_content = doc.page_content
            metadata = doc.metadata
            if page_content not in deduplicated_dict:
                deduplicated_dict[page_content] = (metadata,scores)
        deduplicated_documents = [(Document(page_content=key,metadata=value[0]),value[1]) for key, value in deduplicated_dict.items()]
        return deduplicated_documents
    def similarity_search_with_score_by_vector(
        self,
        embedding: List[float],
        k: int = 4,
        filter: Optional[Dict[str, Any]] = None,
        fetch_k: int = 20,
        **kwargs: Any,
    ) -> List[Tuple[Document, float]]:
        faiss = dependable_faiss_import()
        vector = np.array([embedding], dtype=np.float32)
        if self._normalize_L2:
            faiss.normalize_L2(vector)
        scores, indices = self.index.search(vector, k if filter is None else fetch_k)
        docs = []
        for j, i in enumerate(indices[0]):
            if i == -1:
                # This happens when not enough docs are returned.
                continue
            _id = self.index_to_docstore_id[i]
            doc = self.docstore.search(_id)
            if not isinstance(doc, Document):
                raise ValueError(f"Could not find document for id {_id}, got {doc}")
            if filter is not None:
                filter = {
                    key: [value] if not isinstance(value, list) else value
                    for key, value in filter.items()
                }
                if all(doc.metadata.get(key) in value for key, value in filter.items()):
                    docs.append((doc, scores[0][j]))
            else:
                docs.append((doc, scores[0][j]))
        docs = self._tuple_deduplication(docs)
        score_threshold = kwargs.get("score_threshold")
        if score_threshold is not None:
            cmp = (
                operator.ge
                if self.distance_strategy
                in (DistanceStrategy.MAX_INNER_PRODUCT, DistanceStrategy.JACCARD)
                else operator.le
            )
            docs = [
                (doc, similarity)
                for doc, similarity in docs
                if cmp(similarity, score_threshold)
            ]
        
        if "doc_callback" in kwargs:
            if hasattr(kwargs["doc_callback"], 'after_search'):
                docs = kwargs["doc_callback"].after_search(self.docstore,docs,number=k)
        return docs[:k]
    def max_marginal_relevance_search_by_vector(
        self,
        embedding: List[float],
        k: int = 4,
        fetch_k: int = 20,
        lambda_mult: float = 0.5,
        filter: Optional[Dict[str, Any]] = None,
        **kwargs: Any,
    ) -> List[Document]:
        """Return docs selected using the maximal marginal relevance.

        Maximal marginal relevance optimizes for similarity to query AND diversity
        among selected documents.

        Args:
            embedding: Embedding to look up documents similar to.
            k: Number of Documents to return. Defaults to 4.
            fetch_k: Number of Documents to fetch before filtering to
                     pass to MMR algorithm.
            lambda_mult: Number between 0 and 1 that determines the degree
                        of diversity among the results with 0 corresponding
                        to maximum diversity and 1 to minimum diversity.
                        Defaults to 0.5.
        Returns:
            List of Documents selected by maximal marginal relevance.
        """
        docs_and_scores = self.max_marginal_relevance_search_with_score_by_vector(
            embedding, k=k, fetch_k=fetch_k, lambda_mult=lambda_mult, filter=filter
        )
        docs_and_scores = self._tuple_deduplication(docs_and_scores)
        if "doc_callback" in kwargs:
            if hasattr(kwargs["doc_callback"], 'after_search'):
                docs_and_scores = kwargs["doc_callback"].after_search(self.docstore,docs_and_scores,number=k)
        return [doc for doc, _ in docs_and_scores]
    
def getFAISS(embedding_model_name:str,store_path:str,info:dict = None,index_name:str = "index",is_pgsql:bool = True,reset:bool = False) -> RE_FAISS:
    embeddings = GetEmbding(path=embedding_model_name)
    docstore1:PgSqlDocstore = None
    if is_pgsql:
        if info and "host" in info and "dbname" in info and "username" in info and "password" in info:
            docstore1 = PgSqlDocstore(info,reset=reset)
    else:
        docstore1 = InMemorySecondaryDocstore()
    if not path.exists(store_path):
        os.makedirs(store_path,exist_ok=True)
    if store_path is None or len(store_path) <= 0 or not path.exists(path.join(store_path,index_name+".faiss")) or reset:
        print("create new faiss")
        index = faiss.IndexFlatL2(len(embeddings.embed_documents(["a"])[0])) #根据embeddings向量维度设置
        return RE_FAISS(embedding_function=embeddings.client.encode,index=index,docstore=docstore1,index_to_docstore_id={})
    else:
        print("load_local faiss")
        _faiss = RE_FAISS.load_local(folder_path=store_path,index_name=index_name, embeddings=embeddings)
        if docstore1 and is_pgsql:  #如果外部参数调整，更新docstore
            _faiss.docstore = docstore1
        return _faiss
class VectorStore_FAISS(FAISS):
    def __init__(self,  embedding_model_name:str,store_path:str,index_name:str = "index",info:dict = None, is_pgsql:bool = True,show_number = 5, threshold = 0.8, reset:bool = False,doc_callback:DocumentCallback = DefaultDocumentCallback()):
        self.info = info
        self.embedding_model_name = embedding_model_name
        self.store_path = path.join(store_path,index_name)
        if not path.exists(self.store_path):
            os.makedirs(self.store_path,exist_ok=True)
        self.index_name = index_name
        self.show_number = show_number
        self.search_number = self.show_number*3
        self.threshold = threshold
        self._faiss = getFAISS(self.embedding_model_name,self.store_path,info=info,index_name=self.index_name,is_pgsql=is_pgsql,reset=reset)
        self.doc_callback = doc_callback
    
    def get_text_similarity_with_score(self, text:str,**kwargs):
        score_threshold = (1-self.threshold) * math.sqrt(2)
        docs = self._faiss.similarity_search_with_score(query=text,k=self.search_number,score_threshold=score_threshold,doc_callback=self.doc_callback,**kwargs)
        return [doc for doc, similarity in docs][:self.show_number]
    
    def get_text_similarity(self, text:str,**kwargs):
        docs = self._faiss.similarity_search(query=text,k=self.search_number,doc_callback=self.doc_callback,**kwargs)
        return docs[:self.show_number]
    
    # #去重，并保留metadate
    # def _tuple_deduplication(self, tuple_input:List[Document]) -> List[Document]:
    #     deduplicated_dict = OrderedDict()
    #     for doc in tuple_input:
    #         page_content = doc.page_content
    #         metadata = doc.metadata
    #         if page_content not in deduplicated_dict:
    #             deduplicated_dict[page_content] = metadata
        
    #     deduplicated_documents = [Document(page_content=key,metadata=value) for key, value in deduplicated_dict.items()]
    #     return deduplicated_documents
    
    def _join_document(self, docs:List[Document]) -> str:
        print(docs)
        return "".join([doc.page_content for doc in docs])
        
    def get_local_doc(self, docs:List[Document]):
        ans = []
        for doc in docs:
            ans.append({"page_content":doc.page_content, "page_number":doc.metadata["page_number"], "filename":doc.metadata["filename"]})
        return ans

    # def _join_document_location(self, docs:List[Document]) -> str:

    
    # 持久化到本地
    def _save_local(self):
        self._faiss.save_local(folder_path=self.store_path,index_name=self.index_name)
    # 添加文档
    # Document {
    # page_content 段落
    # metadata {
    #    page 页码    
    #    }    
    # }
    def _add_documents(self, new_docs:List[Document],need_split:bool = True,pattern:str =  r'[？。；\n]'):
        list_of_documents:List[Document] = []
        if self.doc_callback:
            new_docs = self.doc_callback.before_store(self._faiss.docstore,new_docs)
        if need_split:
            for doc in new_docs:
                words_list = re.split(pattern, doc.page_content)
                # 去掉重复项
                words_list = set(words_list)
                words_list = [str(words) for words in words_list]
                for words in words_list:
                    if not words.strip() == '':
                        metadata = copy.deepcopy(doc.metadata)
                        metadata["paragraph"] = doc.page_content
                        list_of_documents.append(Document(page_content=words, metadata=metadata))
        else:
            list_of_documents = new_docs
        self._faiss.add_documents(list_of_documents)
    def _add_documents_from_dir(self,filepaths = [],load_kwargs: Optional[dict] = {"mode":"paged"}):
        self._add_documents(load.loads(filepaths,**load_kwargs))
    def as_retriever(self, **kwargs: Any) -> VectorStoreRetriever:
        """
        Return VectorStoreRetriever initialized from this VectorStore.

        Args:
            search_type (Optional[str]): Defines the type of search that
                the Retriever should perform.
                Can be "similarity" (default), "mmr", or
                "similarity_score_threshold".
            search_kwargs (Optional[Dict]): Keyword arguments to pass to the
                search function. Can include things like:
                    k: Amount of documents to return (Default: 4)
                    score_threshold: Minimum relevance threshold
                        for similarity_score_threshold
                    fetch_k: Amount of documents to pass to MMR algorithm (Default: 20)
                    lambda_mult: Diversity of results returned by MMR;
                        1 for minimum diversity and 0 for maximum. (Default: 0.5)
                    filter: Filter by document metadata

        Returns:
            VectorStoreRetriever: Retriever class for VectorStore.

        Examples:

        .. code-block:: python

            # Retrieve more documents with higher diversity
            # Useful if your dataset has many similar documents
            docsearch.as_retriever(
                search_type="mmr",
                search_kwargs={'k': 6, 'lambda_mult': 0.25}
            )

            # Fetch more documents for the MMR algorithm to consider
            # But only return the top 5
            docsearch.as_retriever(
                search_type="mmr",
                search_kwargs={'k': 5, 'fetch_k': 50}
            )

            # Only retrieve documents that have a relevance score
            # Above a certain threshold
            docsearch.as_retriever(
                search_type="similarity_score_threshold",
                search_kwargs={'score_threshold': 0.8}
            )

            # Only get the single most similar document from the dataset
            docsearch.as_retriever(search_kwargs={'k': 1})

            # Use a filter to only retrieve documents from a specific paper
            docsearch.as_retriever(
                search_kwargs={'filter': {'paper_title':'GPT-4 Technical Report'}}
            )
    """
        if not kwargs or "similarity_score_threshold" != kwargs["search_type"]:
            default_kwargs = {'k': self.show_number}
            if "search_kwargs" in kwargs:
                default_kwargs.update(kwargs["search_kwargs"])
            kwargs["search_kwargs"] = default_kwargs
        elif "similarity_score_threshold" == kwargs["search_type"]:
            default_kwargs = {'score_threshold': self.threshold,'k': self.show_number}
            if "search_kwargs" in kwargs: 
                default_kwargs.update(kwargs["search_kwargs"])
            kwargs["search_kwargs"] = default_kwargs
        kwargs["search_kwargs"]["doc_callback"]=self.doc_callback  
        tags = kwargs.pop("tags", None) or []
        tags.extend(self._faiss._get_retriever_tags())
        print(kwargs)
        return VectorStoreRetriever_FAISS(vectorstore=self._faiss, **kwargs, tags=tags)


class VectorStoreRetriever_FAISS(VectorStoreRetriever):
    search_k = 5
    def __init__(self,**kwargs):
        super().__init__(**kwargs)
        if "k" in self.search_kwargs:
            self.search_k=self.search_kwargs["k"]
            self.search_kwargs["k"]=self.search_k*2
    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> List[Document]:
        docs = super()._get_relevant_documents(query=query,run_manager=run_manager)
        return docs[:self.search_k]
    async def _aget_relevant_documents(
        self, query: str, *, run_manager: AsyncCallbackManagerForRetrieverRun
    ) -> List[Document]:
        docs = super()._aget_relevant_documents(query=query,run_manager=run_manager)
        return docs[:self.search_k]