pinecone_index.py

import os
import pinecone
import time
import yaml

import pandas as pd

from langchain.document_loaders import DataFrameLoader
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores.pinecone import Pinecone
from typing import List

from dotenv import load_dotenv
from pathlib import Path


class PinceconeIndex:
    def __init__(self, index_name: str, model_name: str):
        self.index_name = index_name
        self._embeddingModel = HuggingFaceEmbeddings(model_name=model_name)

    def connect_index(self, embedding_dimension: int, delete_existing: bool = False):
        index_name = self.index_name

        # load pinecone env variables within Google Colab
        if (not os.getenv("PINECONE_KEY")) or (not os.getenv("PINECONE_ENV")):
            dotenv_path = Path("/content/gt-policy-bot/config.env")
            load_dotenv(dotenv_path=dotenv_path)

        pinecone.init(
            api_key=os.getenv("PINECONE_KEY"),
            environment=os.getenv("PINECONE_ENV"),
        )

        if index_name in pinecone.list_indexes() and delete_existing:
            pinecone.delete_index(index_name)

        if index_name not in pinecone.list_indexes():
            pinecone.create_index(index_name, dimension=embedding_dimension)

        index = pinecone.Index(index_name)

        pinecone.describe_index(index_name)
        self._index = index

    def upsert_docs(self, df: pd.DataFrame, text_col: str):
        loader = DataFrameLoader(df, page_content_column=text_col)
        docs = loader.load()
        Pinecone.from_documents(docs, self._embeddingModel, index_name=self.index_name)

    def get_embedding_model(self):
        return self._embeddingModel

    def get_index_name(self):
        return self.index_name

    def query(self, query: str, top_k: int = 5) -> List[str]:
        docsearch = Pinecone.from_existing_index(self.index_name, self._embeddingModel)
        res = docsearch.similarity_search(query, k=top_k)

        return [doc.page_content for doc in res]


if __name__ == "__main__":
    config_path = "config.yml"
    with open("config.yml", "r") as file:
        config = yaml.safe_load(file)

    print(config)

    data_path = config["paths"]["data_path"]
    project = config["paths"]["project"]
    format = ".csv"

    index_name = config["pinecone"]["index-name"]
    embedding_model = config["sentence-transformers"]["model-name"]
    embedding_dimension = config["sentence-transformers"]["embedding-dimension"]
    delete_existing = True

    if config["paths"]["chunking"] == "manual":
        print("Using manual chunking")
        file_path_embedding = config["paths"]["manual_chunk_file"]
        df = pd.read_csv(file_path_embedding, header=None, names=["chunks"])
    else:
        print("Using automatic chunking")
        file_path_embedding = config["paths"]["auto_chunk_file"]
        df = pd.read_csv(file_path_embedding, index_col=0)

    print(df)
    start_time = time.time()
    index = PinceconeIndex(index_name, embedding_model)
    index.connect_index(embedding_dimension, delete_existing)
    index.upsert_docs(df, "chunks")
    end_time = time.time()
    print(f"Indexing took {end_time - start_time} seconds")

    index = PinceconeIndex(index_name, embedding_model)
    index.connect_index(embedding_dimension, delete_existing=False)

    query = "When was the student code of conduct last revised?"
    res = index.query(query, top_k=5)

    # assert len(res) == 5
    print(res)