Update app.py

app.py

@@ -1,4 +1,40 @@
-import gradio as gr
+from langchain.text_splitter import CharacterTextSplitter
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
+from langchain import HuggingFacePipeline
+from langchain.chains import RetrievalQA
+from transformers import AutoTokenizer
+import pickle
+import os
+
+with open('shakespeare.pkl', 'rb') as fp:
+    data = pickle.load(fp)
+
+bloomz_tokenizer = AutoTokenizer.from_pretrained('bigscience/bloomz-1b7')
+
+text_splitter = CharacterTextSplitter.from_huggingface_tokenizer(bloomz_tokenizer, chunk_size=100, chunk_overlap=0, separator='\n')
+
+documents = text_splitter.split_documents(data)
+
+embeddings = HuggingFaceEmbeddings()
+
+persist_directory = "vector_db"
+
+vectordb = Chroma.from_documents(documents=documents, embedding=embeddings, persist_directory=persist_directory)
+
+vectordb.persist()
+vectordb = None
+
+vectordb_persist = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+
+llm = HuggingFacePipeline.from_model_id(
+    model_id="bigscience/bloomz-1b7",
+    task="text-generation",
+    model_kwargs={"temperature" : 0, "max_length" : 500})
+
+doc_retriever = vectordb_persist.as_retriever()
+
+shakespeare_qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=doc_retriever)
 
 def make_inference(query):
     inference = shakespeare_qa.run(query)