fix error

confluence_qa.py

@@ -1,6 +1,6 @@
 from langchain.document_loaders import ConfluenceLoader
 from langchain.text_splitter import CharacterTextSplitter, TokenTextSplitter,RecursiveCharacterTextSplitter,SentenceTransformersTokenTextSplitter
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,pipeline,DistilBertTokenizer,DistilBertForQuestionAnswering
 from langchain.llms.huggingface_pipeline import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
 from langchain.chains import RetrievalQA
@@ -12,10 +12,13 @@ class ConfluenceQA:
         self.embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
     
     def define_model(self) -> None:
-        tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
-        model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
-        pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer, max_new_tokens=1024, truncation=True)
-        self.llm = HuggingFacePipeline(pipeline = pipe,model_kwargs={"temperature": 0, "max_length": 1024})
+        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased-distilled-squad')
+        model = DistilBertForQuestionAnswering.from_pretrained('distilbert-base-uncased-distilled-squad')
+
+        # tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
+        # model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
+        pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+        self.llm = HuggingFacePipeline(pipeline = pipe,model_kwargs={"temperature": 0})
 
     def store_in_vector_db(self) -> None:
         persist_directory = self.config.get("persist_directory",None)
@@ -28,17 +31,12 @@ class ConfluenceQA:
             url=confluence_url, username=username, api_key=api_key
         )
         documents = loader.load(include_attachments=include_attachment, limit=100, space_key=space_key)
-        # text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=10)
-        # documents = text_splitter.split_documents(documents)
-        # print(documents)
-        text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=10)
-        documents = text_splitter.split_documents(documents)
-        # text_splitter = SentenceTransformersTokenTextSplitter(chunk_overlap=10)
-        # documents = text_splitter.split_documents(documents)
-        # print(documents)
-        text_splitter = TokenTextSplitter(chunk_size=1000, chunk_overlap=10, encoding_name="cl100k_base")  # This the encoding for text-embedding-ada-002
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=10)
         documents = text_splitter.split_documents(documents)
         self.db = Chroma.from_documents(documents, self.embeddings)
+        question = "How do I make a space public?"
+        searchDocs = self.db.similarity_search(question)
+        print(searchDocs[0].page_content)
 
 
     def retrieve_qa_chain(self) -> None: