fix

confluence_qa.py

@@ -12,12 +12,12 @@ class ConfluenceQA:
         self.embeddings = HuggingFaceEmbeddings(model_name="multi-qa-MiniLM-L6-cos-v1")
     
     def define_model(self) -> None:
-        # tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased-distilled-squad')
-        # model = DistilBertForQuestionAnswering.from_pretrained('distilbert-base-uncased-distilled-squad')
+        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)
+        # tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-large")
+        # model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-large")
+        pipe = pipeline("question-answering", model=model, tokenizer=tokenizer)
         self.llm = HuggingFacePipeline(pipeline = pipe,model_kwargs={"temperature": 0})
 
     def store_in_vector_db(self) -> None:
@@ -33,13 +33,6 @@ class ConfluenceQA:
         documents = loader.load(include_attachments=include_attachment, limit=100, space_key=space_key)
         text_splitter = RecursiveCharacterTextSplitter(chunk_size=1500, chunk_overlap=150)
         documents = text_splitter.split_documents(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
-        # 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)