fix

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,DistilBertTokenizer,DistilBertForQuestionAnswering
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM,pipeline,T5Tokenizer,T5ForConditionalGeneration
 from langchain.llms.huggingface_pipeline import HuggingFacePipeline
 from langchain.prompts import PromptTemplate
 from langchain.chains import RetrievalQA
@@ -12,13 +12,10 @@ 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 = 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})
+        tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-large")
+        model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-large")
+        pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+        self.llm = HuggingFacePipeline(pipeline = pipe,model_kwargs={"temperature": 0.5})
 
     def store_in_vector_db(self) -> None:
         persist_directory = self.config.get("persist_directory",None)
@@ -31,12 +28,11 @@ 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=1500, chunk_overlap=150)
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, 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)
+        # question = "How do I make a space public?"
+        # searchDocs = self.db.similarity_search(question)
 
     def retrieve_qa_chain(self) -> None:
         template = """Use the following pieces of context to answer the question at the end. If you don't know the answer, just say that you don't know, don't try to make up an answer.
@@ -58,4 +54,5 @@ class ConfluenceQA:
 
     def qa_bot(self, query:str):
         result = self.qa.run(query)
+        print(result)
         return result