使用 Postgresql 数据库保存 Embedding ​

import os
import glob
from typing import List
from multiprocessing import Pool
from tqdm import tqdm
from langchain.document_loaders import (
    CSVLoader,
    EverNoteLoader,
    PDFMinerLoader,
    TextLoader,
    UnstructuredEmailLoader,
    UnstructuredEPubLoader,
    UnstructuredHTMLLoader,
    UnstructuredMarkdownLoader,
    UnstructuredODTLoader,
    UnstructuredPowerPointLoader,
    UnstructuredWordDocumentLoader,
)
from langchain.docstore.document import Document


# Map file extensions to document loaders and their arguments
LOADER_MAPPING = {
    ".csv": (CSVLoader, {}),
    ".doc": (UnstructuredWordDocumentLoader, {}),
    ".docx": (UnstructuredWordDocumentLoader, {}),
    ".enex": (EverNoteLoader, {}),
    ".epub": (UnstructuredEPubLoader, {}),
    ".html": (UnstructuredHTMLLoader, {}),
    ".md": (UnstructuredMarkdownLoader, {}),
    ".odt": (UnstructuredODTLoader, {}),
    ".pdf": (PDFMinerLoader, {}),
    ".ppt": (UnstructuredPowerPointLoader, {}),
    ".pptx": (UnstructuredPowerPointLoader, {}),
    ".txt": (TextLoader, {"encoding": "utf8"}),
}

def load_single_document(file_path: str) -> Document:
    ## Find extension of the file
    ext = "." + file_path.rsplit(".", 1)[-1]
    if ext in LOADER_MAPPING:
        # Find the appropriate loader class and arguments
        loader_class, loader_args = LOADER_MAPPING[ext]
        # Invoke the instance of document loader
        loader = loader_class(file_path, **loader_args)
        ## Return the loaded document
        return loader.load()[0]
    raise ValueError(f"Unsupported file extension '{ext}'")

git_dir = "../../../../privateGPT/"
loaded_document = load_single_document(git_dir+'source_documents/state_of_the_union.txt')
print(f'Type of loaded document {type(loaded_document)}')
loaded_document

def load_documents(source_dir: str, ignored_files: List[str] = []) -> List[Document]:
    """
    Loads all documents from the source documents directory, ignoring specified files
    """
    all_files = []
    for ext in LOADER_MAPPING:
        #Find all the files within source documents which matches the extensions in Loader_Mapping file
        all_files.extend(
            glob.glob(os.path.join(source_dir, f"**/*{ext}"), recursive=True)
        )

    ## Filtering files from all_files if its in ignored_files
    filtered_files = [file_path for file_path in all_files if file_path not in ignored_files]

    ## Spinning up resource pool
    with Pool(processes=os.cpu_count()) as pool:
        results = []
        with tqdm(total=len(filtered_files), desc='Loading new documents', ncols=80) as pbar:
            # Load each document from filtered files list using load_single_document function
            for i, doc in enumerate(pool.imap_unordered(load_single_document, filtered_files)):
                results.append(doc)
                pbar.update()

    return results


loaded_documents = load_documents(git_dir+'source_documents')
print(f"Length of loaded documents: {len(loaded_documents)}")
loaded_documents[0]

from langchain.text_splitter import RecursiveCharacterTextSplitter

chunk_size = 500
chunk_overlap = 50
def process_documents(source_dir: str, ignored_files: List[str] = []) -> List[Document]:
    """
    Load documents and split in chunks
    """
    print(f"Loading documents from {source_dir}")
    documents = load_documents(source_dir, ignored_files)
    if not documents:
        print("No new documents to load")
        exit(0)
    print(f"Loaded {len(documents)} new documents from {source_dir}")
    ## Load text splitter
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    ## Split text
    texts = text_splitter.split_documents(documents)
    print(f"Split into {len(texts)} chunks of text (max. {chunk_size} tokens each)")
    return texts
processed_documents = process_documents(git_dir+'source_documents')

def embed_OpenAi(document:Document):
    """
    使用openai 的接口作获取向量信息
    """
    api_key = os.getenv("OPENAI_API_KEY")

    client = OpenAI(api_key=api_key)
    # Your text input
    text = document.page_content
    # Get the embedding
    response = client.embeddings.create(
        input=text,
        model="text-embedding-ada-002"  # Example model
    )
    # print("response",response)
    # The embedding vector
    embedding_vector = response.data[0].embedding
    return embedding_vector
    # You can now use this vector for further processing

from text2vec import SentenceModel

s_model = SentenceModel("shibing624/text2vec-bge-large-chinese")

def embedding(document:Document):
    """
    使用本地模型处理向量信息
    """
    embeddings = s_model.encode(document.page_content, normalize_embeddings=True)
    return embeddings.tolist()

def insert_vector_data(conn_str, table_name, vector_field, filename,content,vector_data):
    """
    Inserts vector data into a PostgreSQL table with a pgvector column.

    Parameters:
    conn_str (str): The connection string for the database.
    table_name (str): The name of the table to insert data into.
    vector_data (list): The vector data to insert (as a list of floats).
    """
    try:
        # Connect to the database using the connection string
        conn = psycopg2.connect(conn_str)
        cur = conn.cursor()
        from psycopg2.extensions import adapt, AsIs
        vector_data_adapted = AsIs(adapt(vector_data))

        # Insert the vector data
        cur.execute(f"INSERT INTO {table_name} (filename,content,{vector_field}) VALUES (%s,%s,%s);", (filename,content,vector_data_adapted))

        # Commit the transaction and close the connection
        conn.commit()
        cur.close()
        conn.close()

        return "Data inserted successfully."