Open Source Embeddings Optimisation and Eval Framework for RAG/LLM Applications. Documentations at https://docs.vectorboard.ai/introduction
"""
This script loads a PDF document using PyPDFLoader, 
creates a GridSearch object from vectorboard.search,
and runs a grid search on the document using the RetrievalQA chain from langchain.chains.
The grid search uses the specified param_grid and eval_queries to find 
the best combination of parameters for the chain.
The results of the grid search are then shared using a Gradio app link.
"""

from vectorboard.search import GridSearch
from langchain.document_loaders import PyPDFLoader
from langchain.chains import RetrievalQA
from langchain.vectorstores import FAISS
from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings

import os
from dotenv import load_dotenv

load_dotenv()
# os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

loader = PyPDFLoader("recycling.pdf")

param_grid = {
    "chunk_size": [500],
    "vector_store": [FAISS],
    "embeddings": [OpenAIEmbeddings(), HuggingFaceEmbeddings()],
}
eval_queries = [
    "What is the trend in the amount of household waste collected and treated in 2022 compared to the previous year?",  # noqa: E501
    "what percentage of waste is recyvled into materials in 2022?",
    "what percentage of waste is recovered into energy in 2022?",
    "what is the total volume of waste treated for energy recovery per person?",
]


def main():
    vectorboard_grid_search = GridSearch(chain=RetrievalQA)
    vectorboard_grid_search.create_experiments(param_grid=param_grid, loader=loader)
    vectorboard_grid_search.run(eval_queries=eval_queries)
    vectorboard_grid_search.results(share=True)


if __name__ == "__main__":
    main()

src/examples/single-pdf/embed-pdf.py
VectorBoard/vectorboard
from time import monotonic
from rich.console import Console

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.llms import OpenAI


class Experiment:
    """
    A class representing an experiment.

    Attributes:
      params (dict): A dictionary containing experiment parameters.
      index (int): An integer representing the index of the experiment.
      documents (list): A list of documents to be used in the experiment.
      vectorstore_list (list): A list of vector stores used in the experiment.
      retrievers_list (list): A list of retrievers used in the experiment.
      chain (object): An object representing the chain used in the experiment.
      query_results (dict): A dictionary containing the results of the queries.
      run_time (float): A float representing the time it took to run the experiment.
      embedding_time (float): A float representing the time it took to embed the documents.

    Methods:
      run(eval_queries): Runs the experiment with the given evaluation queries.
      _evaluate(chain, eval_queries): Evaluates the experiment with the given chain and evaluation queries.
    """

    def __init__(
        self,
        params,
        index=0,
        documents=None,
        vectorstore_list=None,
        retrievers_list=None,
        chain=None,
    ):
        """
        Initializes an Experiment object.

        Args:
          params (dict): A dictionary containing experiment parameters.
          index (int): An integer representing the index of the experiment.
          documents (list): A list of documents to be used in the experiment.
          vectorstore_list (list): A list of vector stores used in the experiment.
          retrievers_list (list): A list of retrievers used in the experiment.
          chain (object): An object representing the chain used in the experiment.
        """
        self.params = params
        self.index = index
        self.console = Console()
        self.documents = documents
        self.vectorstore_list = vectorstore_list
        self.retrievers_list = retrievers_list
        self.chain = chain

        self.query_results = {}

        # time variables
        self.run_time = 0
        self.embedding_time = 0

    def run(self, eval_queries):
        """
        Runs the experiment with the given evaluation queries.

        Args:
          eval_queries (list): A list of evaluation queries.

        Returns:
          A tuple containing the query results, the run time, and the embedding time.
        """
        self.console.log(f"Experiment {self.index} started")
        self.start_time = monotonic()

        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=self.params["chunk_size"], chunk_overlap=0
        )
        texts = text_splitter.split_documents(self.documents)

        self.pre_embedding = monotonic()
        db = self.params["vector_store"].from_documents(
            texts, self.params["embeddings"]
        )
        self.post_embedding = monotonic()
        self.vectorstore_list.append(db)

        retriever = db.as_retriever(search_kwargs={"k": 2})
        self.retrievers_list.append(retriever)

        qa = self.chain.from_chain_type(
            llm=OpenAI(), chain_type="stuff", retriever=retriever
        )

        self.pre_queries = monotonic()

        self._evaluate(qa, eval_queries)

        return self.query_results, self.run_time, self.embedding_time

    def _evaluate(self, chain, eval_queries):
        """
        Evaluates the experiment with the given chain and evaluation queries.

        Args:
          chain (object): An object representing the chain used in the experiment.
          eval_queries (list): A list of evaluation queries.
        """
        self.console.log(f"Evaluating experiment {self.index}")
        for q in eval_queries:
            res = chain.run(q)
            self.query_results[q] = res

        post_queries = monotonic()
        self.console.log(f"Finished experiment {self.index}")

        self.run_time = post_queries - self.start_time
        self.embedding_time = self.post_embedding - self.pre_embedding

src/vectorboard/experiment.py
VectorBoard/vectorboard
import pandas as pd
from itertools import product
from rich.console import Console

import gradio as gr
import matplotlib.pyplot as plt

from vectorboard.experiment import Experiment


class GridSearch:
    """
    A class for performing grid search on a given chain.

    Attributes:
      experiments (list): A list of Experiment objects.
      experiment_results (DataFrame): A DataFrame containing the results of the experiments.
      console (Console): A Rich Console object.
      vectorstore_list (list): A list of vector stores.
      retrievers_list (list): A list of retrievers.
      chain (Chain): A Chain object.

    Methods:
      create_experiments(loader, param_grid): Creates a list of Experiment objects based on the given parameter grid.
      run(eval_queries): Runs the experiments on the given evaluation queries.
      results(gradio_app=True, share=False): Displays the results of the experiments in a Gradio app.
    """

    def __init__(self, chain):
        """
        Initializes a GridSearch object.

        Args:
          chain (Chain): A Chain object.
        """
        self.experiments = []
        self.experiment_results = None
        self.console = Console()
        self.vectorstore_list = []
        self.retrievers_list = []
        self.chain = chain

    def create_experiments(self, param_grid, loader=None, documents=None):
        """
        Creates a list of Experiment objects based on the given parameter grid.

        Args:
          param_grid (dict): A dictionary containing the parameter grid.
          loader (Loader): A Loader object.
          documents (list): A list of documents.

        Raises:
          ValueError: If the parameter grid does not contain 'embeddings', 'vector_store', or 'chunk_size'.
        """
        if "embeddings" not in param_grid.keys():
            raise ValueError("param_grid must contain embeddings")
        elif "vector_store" not in param_grid.keys():
            raise ValueError("param_grid must contain vector_store")
        elif "chunk_size" not in param_grid.keys():
            raise ValueError("param_grid must contain chunk_size")
        else:
            pass

        if documents is not None:
            self.documents = documents
        elif loader is not None:
            self.documents = loader.load()
        else:
            raise ValueError("Either documents or loader must be provided")

        all_params = [
            dict(zip(param_grid.keys(), values))
            for values in product(*param_grid.values())
        ]

        self.experiments_info_df = pd.DataFrame(
            all_params,
            index=[f"Experiment_{i}" for i in range(len(all_params))],
        )

        self.experiments = [
            Experiment(
                params=params,
                index=i + 1,  # to start from 1 instead of 0
                documents=self.documents,
                vectorstore_list=self.vectorstore_list,
                retrievers_list=self.retrievers_list,
                chain=self.chain,
            )
            for i, params in enumerate(all_params)
        ]

    def run(self, eval_queries):
        """
        Runs the experiments on the given evaluation queries.

        Args:
          eval_queries (list): A list of evaluation queries.
        """
        self.experiments_res_df = pd.DataFrame(index=eval_queries)

        for i, experiment in enumerate(self.experiments):
            response, runtime, embedding_time = experiment.run(eval_queries)

            self.experiments_res_df["Experiment_" + str(i)] = response.values()

            self.experiments_info_df.at["Experiment_" + str(i), "run time"] = runtime

            self.experiments_info_df.at[
                "Experiment_" + str(i), "embedding time"
            ] = embedding_time

    def results(self, gradio_app=True, share=False):
        """
        Displays the results of the experiments in a Gradio app.

        Args:
          gradio_app (bool): Whether to display the results in a Gradio app or not.
          share (bool): Whether to allow sharing of the Gradio app or not.
        """
        if gradio_app:

            def plot_fn():
                plt.figure(figsize=(50, 35))
                plt.style.use("seaborn-v0_8-poster")
                plt.style.use("seaborn-v0_8-deep")
                ax = self.experiments_info_df[["run time", "embedding time"]].plot.bar(
                    subplots=True, figsize=(40, 25), ylabel="Time (s)"
                )

                for axis in ax:
                    for p in axis.patches:
                        axis.annotate(
                            f"{p.get_height():.2f}",
                            (p.get_x() + p.get_width() / 2.0, p.get_height() / 2),
                            ha="center",
                            va="center",
                            color="white",
                            rotation=0,
                        )

                plt.xticks(rotation=30, horizontalalignment="center")
                plt.axhline(0, color="k")
                return plt

            with gr.Blocks(theme=gr.themes.Soft()) as demo:
                gr.Markdown(
                    """
            # 💛 Vectorboard Grid Search  
            ## Experiment results
          """
                )
                experiments = gr.Dataframe(self.experiments_info_df)
                queries = gr.Dataframe(
                    pd.DataFrame(self.experiments_res_df.index, columns=["Queries"])
                )
                result_table = gr.Dataframe(
                    self.experiments_res_df, show_label=True, label="Result table"
                )
                plot = gr.Plot(label="Plot")
                demo.load(plot_fn, [], [plot])

            demo.launch(share=True)
        else:
            pass

src/vectorboard/search.py
VectorBoard/vectorboard