Update main.py

main.py

@@ -1,50 +1,52 @@
-from langchain.chains import RetrievalQA, ConversationalRetrievalChain
-from langchain.vectorstores import Chroma
-from langchain.text_splitter import CharacterTextSplitter
-from langchain.document_loaders import DirectoryLoader, TextLoader,PyPDFLoader
-from transformers import pipeline, AutoModelForCausalLM
+from langchain_community.document_loaders import TextLoader
+from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+from langchain_community.vectorstores import FAISS
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.llms import HuggingFacePipeline
-from langchain.embeddings import HuggingFaceInstructEmbeddings
-import gradio as gr
-from InstructorEmbedding import INSTRUCTOR
-import torch
+from transformers import pipeline
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from langchain.chains import RetrievalQA
+import torch
+loader = TextLoader("info.txt")
+docs = loader.load()
+text_splitter = RecursiveCharacterTextSplitter()
+# text_splitter = RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
+documents = text_splitter.split_documents(docs)
+
+huggingface_embeddings = HuggingFaceBgeEmbeddings(
+    model_name="BAAI/bge-small-en-v1.5",
+    model_kwargs={'device':'cpu'},
+    encode_kwargs={'normalize_embeddings': True}
+)
+
+vector = FAISS.from_documents(documents, huggingface_embeddings)
+retriever = vector.as_retriever()
 
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
+model_name = "facebook/bart-base"
 
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
-tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-base")
-model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
 
 pipe = pipeline(
     "text2text-generation",
     model=model,
     tokenizer=tokenizer,
-    max_length=200,
-    temperature=0.8,
-    top_p=0.95,
+    max_length=300,
+    temperature=0.9,
+    top_p=0.9,
     repetition_penalty=1.15,
     do_sample=True
-)
 
+)
 local_llm = HuggingFacePipeline(pipeline=pipe)
-loader = PyPDFLoader('bipolar.pdf')
-# loader = TextLoader('info.txt')
-document = loader.load()
-text_spliter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
-texts = text_spliter.split_documents(document)
-embedding = HuggingFaceInstructEmbeddings()
-docsearch = Chroma.from_documents(texts, embedding, persist_directory='db')
-
-retriever = docsearch.as_retriever(search_kwargs={"k": 3})
-qa_chain = RetrievalQA.from_chain_type(llm=local_llm,
-                                       chain_type="stuff",
-                                       retriever=retriever,
-                                       return_source_documents=True)
+qa_chain =  RetrievalQA.from_llm(llm=local_llm, retriever=retriever)
+
+
 
 def gradinterface(query,history):
     result = qa_chain({'query': query})
-    return result['result']
+    return result
 
 
 demo = gr.ChatInterface(fn=gradinterface, title='OUR_OWN_BOT')