move app to pdf_qa and create app

README.md

@@ -12,3 +12,7 @@ pinned: false
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
 
 inspired by source https://www.shakudo.io/blog/build-pdf-bot-open-source-llms
+
+
+# Deployed on hugging face 
+https://huggingface.co/spaces/spoggy/streamlit_pdf_qna_open_models

app.py

@@ -1,246 +1,91 @@
-#from langchain.document_loaders.pdf import PDFPlumberLoader
-from langchain.document_loaders import PyPDFLoader
-from langchain.text_splitter import CharacterTextSplitter, TokenTextSplitter
-from transformers import pipeline
-from langchain.prompts import PromptTemplate
-from langchain.chat_models import ChatOpenAI
-from langchain.vectorstores import Chroma
-from langchain.chains import RetrievalQA
-from langchain.embeddings import HuggingFaceInstructEmbeddings, HuggingFaceEmbeddings
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.llms import OpenAI, HuggingFacePipeline
+import streamlit as st
+from pdf_qa import PdfQA
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+import time
+import shutil
 from constants import *
-from transformers import AutoTokenizer
-import torch
-import os
-import re
-from pprint import pprint
 
-class PdfQA:
-    def __init__(self,config:dict = {}):
-        self.config = config
-        self.embedding = None
-        self.vectordb = None
-        self.llm = None
-        self.qa = None
-        self.retriever = None
 
-    # The following class methods are useful to create global GPU model instances
-    # This way we don't need to reload models in an interactive app,
-    # and the same model instance can be used across multiple user sessions
-    @classmethod
-    def create_instructor_xl(cls):
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        return HuggingFaceInstructEmbeddings(model_name=EMB_INSTRUCTOR_XL, model_kwargs={"device": device})
-    
-    @classmethod
-    def create_sbert_mpnet(cls):
-        device = "cuda" if torch.cuda.is_available() else "cpu"
-        return HuggingFaceEmbeddings(model_name=EMB_SBERT_MPNET_BASE, model_kwargs={"device": device})    
-    
-    @classmethod
-    def create_flan_t5_xxl(cls, load_in_8bit=False):
-        # Local flan-t5-xxl with 8-bit quantization for inference
-        # Wrap it in HF pipeline for use with LangChain
-        return pipeline(
-            task="text2text-generation",
-            model="google/flan-t5-xxl",
-            max_new_tokens=200,
-            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
-        )
-    @classmethod
-    def create_flan_t5_xl(cls, load_in_8bit=False):
-        return pipeline(
-            task="text2text-generation",
-            model="google/flan-t5-xl",
-            max_new_tokens=200,
-            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
-        )
-    
-    @classmethod
-    def create_flan_t5_small(cls, load_in_8bit=False):
-        # Local flan-t5-small for inference
-        # Wrap it in HF pipeline for use with LangChain
-        model="google/flan-t5-small"
-        tokenizer = AutoTokenizer.from_pretrained(model)
-        return pipeline(
-            task="text2text-generation",
-            model=model,
-            tokenizer = tokenizer,
-            max_new_tokens=100,
-            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
-        )
-    @classmethod
-    def create_flan_t5_base(cls, load_in_8bit=False):
-        # Wrap it in HF pipeline for use with LangChain
-        model="google/flan-t5-base"
-        tokenizer = AutoTokenizer.from_pretrained(model)
-        return pipeline(
-            task="text2text-generation",
-            model=model,
-            tokenizer = tokenizer,
-            max_new_tokens=100,
-            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
-        )
-    @classmethod
-    def create_flan_t5_large(cls, load_in_8bit=False):
-        # Wrap it in HF pipeline for use with LangChain
-        model="google/flan-t5-large"
-        tokenizer = AutoTokenizer.from_pretrained(model)
-        return pipeline(
-            task="text2text-generation",
-            model=model,
-            tokenizer = tokenizer,
-            max_new_tokens=100,
-            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
-        )
-    @classmethod
-    def create_fastchat_t5_xl(cls, load_in_8bit=False):
-        return pipeline(
-            task="text2text-generation",
-            model = "lmsys/fastchat-t5-3b-v1.0",
-            max_new_tokens=100,
-            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
-        )
-    
-    @classmethod
-    def create_falcon_instruct_small(cls, load_in_8bit=False):
-        model = "tiiuae/falcon-7b-instruct"
 
-        tokenizer = AutoTokenizer.from_pretrained(model)
-        hf_pipeline = pipeline(
-                task="text-generation",
-                model = model,
-                tokenizer = tokenizer,
-                trust_remote_code = True,
-                max_new_tokens=100,
-                model_kwargs={
-                    "device_map": "auto", 
-                    "load_in_8bit": load_in_8bit, 
-                    "max_length": 512, 
-                    "temperature": 0.01,
-                    "torch_dtype":torch.bfloat16,
-                    }
-            )
-        return hf_pipeline
-    
-    def init_embeddings(self) -> None:
-        # OpenAI ada embeddings API
-        if self.config["embedding"] == EMB_OPENAI_ADA:
-            self.embedding = OpenAIEmbeddings()
-        elif self.config["embedding"] == EMB_INSTRUCTOR_XL:
-            # Local INSTRUCTOR-XL embeddings
-            if self.embedding is None:
-                self.embedding = PdfQA.create_instructor_xl()
-        elif self.config["embedding"] == EMB_SBERT_MPNET_BASE:
-            ## this is for SBERT
-            if self.embedding is None:
-                self.embedding = PdfQA.create_sbert_mpnet()
-        else:
-            self.embedding = None ## DuckDb uses sbert embeddings
-            # raise ValueError("Invalid config")
+# Streamlit app code
+st.set_page_config(
+    page_title='Q&A Bot for PDF',
+    page_icon='🔖',
+    layout='wide',
+    initial_sidebar_state='auto',
+)
+
+
+if "pdf_qa_model" not in st.session_state:
+    st.session_state["pdf_qa_model"]:PdfQA = PdfQA() ## Intialisation
 
-    def init_models(self) -> None:
-        """ Initialize LLM models based on config """
-        load_in_8bit = self.config.get("load_in_8bit",False)
-        # OpenAI GPT 3.5 API
-        if self.config["llm"] == LLM_OPENAI_GPT35:
-            # OpenAI GPT 3.5 API
-            pass
-        elif self.config["llm"] == LLM_FLAN_T5_SMALL:
-            if self.llm is None:
-                self.llm = PdfQA.create_flan_t5_small(load_in_8bit=load_in_8bit)
-        elif self.config["llm"] == LLM_FLAN_T5_BASE:
-            if self.llm is None:
-                self.llm = PdfQA.create_flan_t5_base(load_in_8bit=load_in_8bit)
-        elif self.config["llm"] == LLM_FLAN_T5_LARGE:
-            if self.llm is None:
-                self.llm = PdfQA.create_flan_t5_large(load_in_8bit=load_in_8bit)
-        elif self.config["llm"] == LLM_FLAN_T5_XL:
-            if self.llm is None:
-                self.llm = PdfQA.create_flan_t5_xl(load_in_8bit=load_in_8bit)
-        elif self.config["llm"] == LLM_FLAN_T5_XXL:
-            if self.llm is None:
-                self.llm = PdfQA.create_flan_t5_xxl(load_in_8bit=load_in_8bit)
-        elif self.config["llm"] == LLM_FASTCHAT_T5_XL:
-            if self.llm is None:
-                self.llm = PdfQA.create_fastchat_t5_xl(load_in_8bit=load_in_8bit)
-        elif self.config["llm"] == LLM_FALCON_SMALL:
-            if self.llm is None:
-                self.llm = PdfQA.create_falcon_instruct_small(load_in_8bit=load_in_8bit)
-        
-        else:
-            raise ValueError("Invalid config")        
-    def vector_db_pdf(self) -> None:
-        """
-        creates vector db for the embeddings and persists them or loads a vector db from the persist directory
-        """
-        pdf_path = self.config.get("pdf_path",None)
-        persist_directory = self.config.get("persist_directory",None)
-        if persist_directory and os.path.exists(persist_directory):
-            ## Load from the persist db
-            self.vectordb = Chroma(persist_directory=persist_directory, embedding_function=self.embedding)
-        elif pdf_path and os.path.exists(pdf_path):
-            ## 1. Extract the documents
-            loader = PyPDFLoader(pdf_path)
-            documents = loader.load()
-            ## 2. Split the texts
-            text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
-            texts = text_splitter.split_documents(documents)
-            # text_splitter = TokenTextSplitter(chunk_size=100, chunk_overlap=10, encoding_name="cl100k_base")  # This the encoding for text-embedding-ada-002
-            text_splitter = TokenTextSplitter(chunk_size=100, chunk_overlap=10)  # This the encoding for text-embedding-ada-002
-            texts = text_splitter.split_documents(texts)
+## To cache resource across multiple session 
+@st.cache_resource
+def load_llm(llm,load_in_8bit):
 
-            ## 3. Create Embeddings and add to chroma store
-            ##TODO: Validate if self.embedding is not None
-            self.vectordb = Chroma.from_documents(documents=texts, embedding=self.embedding, persist_directory=persist_directory)
-        else:
-            raise ValueError("NO PDF found")
+    if llm == LLM_OPENAI_GPT35:
+        pass
+    elif llm == LLM_FLAN_T5_SMALL:
+        return PdfQA.create_flan_t5_small(load_in_8bit)
+    elif llm == LLM_FLAN_T5_BASE:
+        return PdfQA.create_flan_t5_base(load_in_8bit)
+    elif llm == LLM_FLAN_T5_LARGE:
+        return PdfQA.create_flan_t5_large(load_in_8bit)
+    elif llm == LLM_FASTCHAT_T5_XL:
+        return PdfQA.create_fastchat_t5_xl(load_in_8bit)
+    elif llm == LLM_FALCON_SMALL:
+        return PdfQA.create_falcon_instruct_small(load_in_8bit)
+    else:
+        raise ValueError("Invalid LLM setting")
 
-    def retreival_qa_chain(self):
-        """
-        Creates retrieval qa chain using vectordb as retrivar and LLM to complete the prompt
-        """
-        ##TODO: Use custom prompt
-        print("one", self)
-        pprint(vars(self))
-        self.retriever = self.vectordb.as_retriever(search_kwargs={"k":3})
-        print("two")
-        
-        if self.config["llm"] == LLM_OPENAI_GPT35:
-          # Use ChatGPT API
-          self.qa = RetrievalQA.from_chain_type(llm=OpenAI(model_name=LLM_OPENAI_GPT35, temperature=0.), chain_type="stuff",\
-                                      retriever=self.vectordb.as_retriever(search_kwargs={"k":3}))
-        else:
-            hf_llm = HuggingFacePipeline(pipeline=self.llm,model_id=self.config["llm"])
+## To cache resource across multiple session
+@st.cache_resource
+def load_emb(emb):
+    if emb == EMB_INSTRUCTOR_XL:
+        return PdfQA.create_instructor_xl()
+    elif emb == EMB_SBERT_MPNET_BASE:
+        return PdfQA.create_sbert_mpnet()
+    elif emb == EMB_SBERT_MINILM:
+        pass ##ChromaDB takes care
+    else:
+        raise ValueError("Invalid embedding setting")
+
+
+
+st.title("PDF Q&A (Self hosted LLMs)")
+
+with st.sidebar:
+    emb = st.radio("**Select Embedding Model**", [EMB_INSTRUCTOR_XL, EMB_SBERT_MPNET_BASE,EMB_SBERT_MINILM],index=1)
+    llm = st.radio("**Select LLM Model**", [LLM_FASTCHAT_T5_XL, LLM_FLAN_T5_SMALL,LLM_FLAN_T5_BASE,LLM_FLAN_T5_LARGE,LLM_FLAN_T5_XL,LLM_FALCON_SMALL],index=2)
+    load_in_8bit = st.radio("**Load 8 bit**", [True, False],index=1)
+    pdf_file = st.file_uploader("**Upload PDF**", type="pdf")
+
+    
+    if st.button("Submit") and pdf_file is not None:
+        with st.spinner(text="Uploading PDF and Generating Embeddings.."):
+            with NamedTemporaryFile(delete=False, suffix='.pdf') as tmp:
+                shutil.copyfileobj(pdf_file, tmp)
+                tmp_path = Path(tmp.name)
+                st.session_state["pdf_qa_model"].config = {
+                    "pdf_path": str(tmp_path),
+                    "embedding": emb,
+                    "llm": llm,
+                    "load_in_8bit": load_in_8bit
+                }
+                st.session_state["pdf_qa_model"].embedding = load_emb(emb)
+                st.session_state["pdf_qa_model"].llm = load_llm(llm,load_in_8bit)        
+                st.session_state["pdf_qa_model"].init_embeddings()
+                st.session_state["pdf_qa_model"].init_models()
+                st.session_state["pdf_qa_model"].vector_db_pdf()
+                st.sidebar.success("PDF uploaded successfully")
 
-            self.qa = RetrievalQA.from_chain_type(llm=hf_llm, chain_type="stuff",retriever=self.retriever)
-            if self.config["llm"] == LLM_FLAN_T5_SMALL or self.config["llm"] == LLM_FLAN_T5_BASE or self.config["llm"] == LLM_FLAN_T5_LARGE:
-                question_t5_template = """
-                context: {context}
-                question: {question}
-                answer: 
-                """
-                QUESTION_T5_PROMPT = PromptTemplate(
-                    template=question_t5_template, input_variables=["context", "question"]
-                )
-                self.qa.combine_documents_chain.llm_chain.prompt = QUESTION_T5_PROMPT
-            self.qa.combine_documents_chain.verbose = True
-            self.qa.return_source_documents = True
-    def answer_query(self,question:str) ->str:
-        """
-        Answer the question
-        """
+question = st.text_input('Ask a question', 'What is this document?')
 
-        answer_dict = self.qa({"query":question,})
-        print(answer_dict)
-        answer = answer_dict["result"]
-        if self.config["llm"] == LLM_FASTCHAT_T5_XL:
-            answer = self._clean_fastchat_t5_output(answer)
-        return answer
-    def _clean_fastchat_t5_output(self, answer: str) -> str:
-        # Remove <pad> tags, double spaces, trailing newline
-        answer = re.sub(r"<pad>\s+", "", answer)
-        answer = re.sub(r"  ", " ", answer)
-        answer = re.sub(r"\n$", "", answer)
-        return answer
+if st.button("Answer"):
+    try:
+        st.session_state["pdf_qa_model"].retreival_qa_chain()
+        answer = st.session_state["pdf_qa_model"].answer_query(question)
+        st.write(f"{answer}")
+    except Exception as e:
+        st.error(f"Error answering the question: {str(e)}")

constants.py

@@ -0,0 +1,15 @@
+# Constants
+EMB_OPENAI_ADA = "text-embedding-ada-002"
+EMB_INSTRUCTOR_XL = "hkunlp/instructor-xl"
+EMB_SBERT_MPNET_BASE = "sentence-transformers/all-mpnet-base-v2" # Chroma takes care if embeddings are None
+EMB_SBERT_MINILM = "sentence-transformers/all-MiniLM-L6-v2" # Chroma takes care if embeddings are None
+
+
+LLM_OPENAI_GPT35 = "gpt-3.5-turbo"
+LLM_FLAN_T5_XXL = "google/flan-t5-xxl"
+LLM_FLAN_T5_XL = "google/flan-t5-xl"
+LLM_FASTCHAT_T5_XL = "lmsys/fastchat-t5-3b-v1.0"
+LLM_FLAN_T5_SMALL = "google/flan-t5-small"
+LLM_FLAN_T5_BASE = "google/flan-t5-base"
+LLM_FLAN_T5_LARGE = "google/flan-t5-large"
+LLM_FALCON_SMALL = "tiiuae/falcon-7b-instruct"

pdf_qa.py

@@ -0,0 +1,246 @@
+#from langchain.document_loaders.pdf import PDFPlumberLoader
+from langchain.document_loaders import PyPDFLoader
+from langchain.text_splitter import CharacterTextSplitter, TokenTextSplitter
+from transformers import pipeline
+from langchain.prompts import PromptTemplate
+from langchain.chat_models import ChatOpenAI
+from langchain.vectorstores import Chroma
+from langchain.chains import RetrievalQA
+from langchain.embeddings import HuggingFaceInstructEmbeddings, HuggingFaceEmbeddings
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.llms import OpenAI, HuggingFacePipeline
+from constants import *
+from transformers import AutoTokenizer
+import torch
+import os
+import re
+from pprint import pprint
+
+class PdfQA:
+    def __init__(self,config:dict = {}):
+        self.config = config
+        self.embedding = None
+        self.vectordb = None
+        self.llm = None
+        self.qa = None
+        self.retriever = None
+
+    # The following class methods are useful to create global GPU model instances
+    # This way we don't need to reload models in an interactive app,
+    # and the same model instance can be used across multiple user sessions
+    @classmethod
+    def create_instructor_xl(cls):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        return HuggingFaceInstructEmbeddings(model_name=EMB_INSTRUCTOR_XL, model_kwargs={"device": device})
+    
+    @classmethod
+    def create_sbert_mpnet(cls):
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        return HuggingFaceEmbeddings(model_name=EMB_SBERT_MPNET_BASE, model_kwargs={"device": device})    
+    
+    @classmethod
+    def create_flan_t5_xxl(cls, load_in_8bit=False):
+        # Local flan-t5-xxl with 8-bit quantization for inference
+        # Wrap it in HF pipeline for use with LangChain
+        return pipeline(
+            task="text2text-generation",
+            model="google/flan-t5-xxl",
+            max_new_tokens=200,
+            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
+        )
+    @classmethod
+    def create_flan_t5_xl(cls, load_in_8bit=False):
+        return pipeline(
+            task="text2text-generation",
+            model="google/flan-t5-xl",
+            max_new_tokens=200,
+            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
+        )
+    
+    @classmethod
+    def create_flan_t5_small(cls, load_in_8bit=False):
+        # Local flan-t5-small for inference
+        # Wrap it in HF pipeline for use with LangChain
+        model="google/flan-t5-small"
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        return pipeline(
+            task="text2text-generation",
+            model=model,
+            tokenizer = tokenizer,
+            max_new_tokens=100,
+            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
+        )
+    @classmethod
+    def create_flan_t5_base(cls, load_in_8bit=False):
+        # Wrap it in HF pipeline for use with LangChain
+        model="google/flan-t5-base"
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        return pipeline(
+            task="text2text-generation",
+            model=model,
+            tokenizer = tokenizer,
+            max_new_tokens=100,
+            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
+        )
+    @classmethod
+    def create_flan_t5_large(cls, load_in_8bit=False):
+        # Wrap it in HF pipeline for use with LangChain
+        model="google/flan-t5-large"
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        return pipeline(
+            task="text2text-generation",
+            model=model,
+            tokenizer = tokenizer,
+            max_new_tokens=100,
+            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
+        )
+    @classmethod
+    def create_fastchat_t5_xl(cls, load_in_8bit=False):
+        return pipeline(
+            task="text2text-generation",
+            model = "lmsys/fastchat-t5-3b-v1.0",
+            max_new_tokens=100,
+            model_kwargs={"device_map": "auto", "load_in_8bit": load_in_8bit, "max_length": 512, "temperature": 0.}
+        )
+    
+    @classmethod
+    def create_falcon_instruct_small(cls, load_in_8bit=False):
+        model = "tiiuae/falcon-7b-instruct"
+
+        tokenizer = AutoTokenizer.from_pretrained(model)
+        hf_pipeline = pipeline(
+                task="text-generation",
+                model = model,
+                tokenizer = tokenizer,
+                trust_remote_code = True,
+                max_new_tokens=100,
+                model_kwargs={
+                    "device_map": "auto", 
+                    "load_in_8bit": load_in_8bit, 
+                    "max_length": 512, 
+                    "temperature": 0.01,
+                    "torch_dtype":torch.bfloat16,
+                    }
+            )
+        return hf_pipeline
+    
+    def init_embeddings(self) -> None:
+        # OpenAI ada embeddings API
+        if self.config["embedding"] == EMB_OPENAI_ADA:
+            self.embedding = OpenAIEmbeddings()
+        elif self.config["embedding"] == EMB_INSTRUCTOR_XL:
+            # Local INSTRUCTOR-XL embeddings
+            if self.embedding is None:
+                self.embedding = PdfQA.create_instructor_xl()
+        elif self.config["embedding"] == EMB_SBERT_MPNET_BASE:
+            ## this is for SBERT
+            if self.embedding is None:
+                self.embedding = PdfQA.create_sbert_mpnet()
+        else:
+            self.embedding = None ## DuckDb uses sbert embeddings
+            # raise ValueError("Invalid config")
+
+    def init_models(self) -> None:
+        """ Initialize LLM models based on config """
+        load_in_8bit = self.config.get("load_in_8bit",False)
+        # OpenAI GPT 3.5 API
+        if self.config["llm"] == LLM_OPENAI_GPT35:
+            # OpenAI GPT 3.5 API
+            pass
+        elif self.config["llm"] == LLM_FLAN_T5_SMALL:
+            if self.llm is None:
+                self.llm = PdfQA.create_flan_t5_small(load_in_8bit=load_in_8bit)
+        elif self.config["llm"] == LLM_FLAN_T5_BASE:
+            if self.llm is None:
+                self.llm = PdfQA.create_flan_t5_base(load_in_8bit=load_in_8bit)
+        elif self.config["llm"] == LLM_FLAN_T5_LARGE:
+            if self.llm is None:
+                self.llm = PdfQA.create_flan_t5_large(load_in_8bit=load_in_8bit)
+        elif self.config["llm"] == LLM_FLAN_T5_XL:
+            if self.llm is None:
+                self.llm = PdfQA.create_flan_t5_xl(load_in_8bit=load_in_8bit)
+        elif self.config["llm"] == LLM_FLAN_T5_XXL:
+            if self.llm is None:
+                self.llm = PdfQA.create_flan_t5_xxl(load_in_8bit=load_in_8bit)
+        elif self.config["llm"] == LLM_FASTCHAT_T5_XL:
+            if self.llm is None:
+                self.llm = PdfQA.create_fastchat_t5_xl(load_in_8bit=load_in_8bit)
+        elif self.config["llm"] == LLM_FALCON_SMALL:
+            if self.llm is None:
+                self.llm = PdfQA.create_falcon_instruct_small(load_in_8bit=load_in_8bit)
+        
+        else:
+            raise ValueError("Invalid config")        
+    def vector_db_pdf(self) -> None:
+        """
+        creates vector db for the embeddings and persists them or loads a vector db from the persist directory
+        """
+        pdf_path = self.config.get("pdf_path",None)
+        persist_directory = self.config.get("persist_directory",None)
+        if persist_directory and os.path.exists(persist_directory):
+            ## Load from the persist db
+            self.vectordb = Chroma(persist_directory=persist_directory, embedding_function=self.embedding)
+        elif pdf_path and os.path.exists(pdf_path):
+            ## 1. Extract the documents
+            loader = PyPDFLoader(pdf_path)
+            documents = loader.load()
+            ## 2. Split the texts
+            text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=0)
+            texts = text_splitter.split_documents(documents)
+            # text_splitter = TokenTextSplitter(chunk_size=100, chunk_overlap=10, encoding_name="cl100k_base")  # This the encoding for text-embedding-ada-002
+            text_splitter = TokenTextSplitter(chunk_size=100, chunk_overlap=10)  # This the encoding for text-embedding-ada-002
+            texts = text_splitter.split_documents(texts)
+
+            ## 3. Create Embeddings and add to chroma store
+            ##TODO: Validate if self.embedding is not None
+            self.vectordb = Chroma.from_documents(documents=texts, embedding=self.embedding, persist_directory=persist_directory)
+        else:
+            raise ValueError("NO PDF found")
+
+    def retreival_qa_chain(self):
+        """
+        Creates retrieval qa chain using vectordb as retrivar and LLM to complete the prompt
+        """
+        ##TODO: Use custom prompt
+        print("one", self)
+        pprint(vars(self))
+        self.retriever = self.vectordb.as_retriever(search_kwargs={"k":3})
+        print("two")
+        
+        if self.config["llm"] == LLM_OPENAI_GPT35:
+          # Use ChatGPT API
+          self.qa = RetrievalQA.from_chain_type(llm=OpenAI(model_name=LLM_OPENAI_GPT35, temperature=0.), chain_type="stuff",\
+                                      retriever=self.vectordb.as_retriever(search_kwargs={"k":3}))
+        else:
+            hf_llm = HuggingFacePipeline(pipeline=self.llm,model_id=self.config["llm"])
+
+            self.qa = RetrievalQA.from_chain_type(llm=hf_llm, chain_type="stuff",retriever=self.retriever)
+            if self.config["llm"] == LLM_FLAN_T5_SMALL or self.config["llm"] == LLM_FLAN_T5_BASE or self.config["llm"] == LLM_FLAN_T5_LARGE:
+                question_t5_template = """
+                context: {context}
+                question: {question}
+                answer: 
+                """
+                QUESTION_T5_PROMPT = PromptTemplate(
+                    template=question_t5_template, input_variables=["context", "question"]
+                )
+                self.qa.combine_documents_chain.llm_chain.prompt = QUESTION_T5_PROMPT
+            self.qa.combine_documents_chain.verbose = True
+            self.qa.return_source_documents = True
+    def answer_query(self,question:str) ->str:
+        """
+        Answer the question
+        """
+
+        answer_dict = self.qa({"query":question,})
+        print(answer_dict)
+        answer = answer_dict["result"]
+        if self.config["llm"] == LLM_FASTCHAT_T5_XL:
+            answer = self._clean_fastchat_t5_output(answer)
+        return answer
+    def _clean_fastchat_t5_output(self, answer: str) -> str:
+        # Remove <pad> tags, double spaces, trailing newline
+        answer = re.sub(r"<pad>\s+", "", answer)
+        answer = re.sub(r"  ", " ", answer)
+        answer = re.sub(r"\n$", "", answer)
+        return answer