gkrishnan commited on
Commit
2063044
1 Parent(s): c5fa6de

Create app.py

Browse files

Shakespeare app

Files changed (1) hide show
  1. app.py +53 -0
app.py ADDED
@@ -0,0 +1,53 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain.text_splitter import CharacterTextSplitter
2
+ from langchain.embeddings import HuggingFaceEmbeddings
3
+ from langchain.vectorstores import Chroma
4
+ from langchain import HuggingFacePipeline
5
+ from langchain.chains import RetrievalQA
6
+ from transformers import AutoTokenizer
7
+ import pickle
8
+ import os
9
+
10
+ with open('shakespeare.pkl', 'rb') as fp:
11
+ data = pickle.load(fp)
12
+
13
+ bloomz_tokenizer = AutoTokenizer.from_pretrained('bigscience/bloomz-1b7')
14
+
15
+ text_splitter = CharacterTextSplitter.from_huggingface_tokenizer(bloomz_tokenizer, chunk_size=100, chunk_overlap=0, separator='\n')
16
+
17
+ documents = text_splitter.split_documents(data)
18
+
19
+ embeddings = HuggingFaceEmbeddings()
20
+
21
+ persist_directory = "vector_db"
22
+
23
+ vectordb = Chroma.from_documents(documents=documents, embedding=embeddings, persist_directory=persist_directory)
24
+
25
+ vectordb.persist()
26
+ vectordb = None
27
+
28
+ vectordb_persist = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
29
+
30
+ llm = HuggingFacePipeline.from_model_id(
31
+ model_id="bigscience/bloomz-1b7",
32
+ task="text-generation",
33
+ model_kwargs={"temperature" : 0, "max_length" : 500})
34
+
35
+ doc_retriever = vectordb_persist.as_retriever()
36
+
37
+ shakespeare_qa = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff", retriever=doc_retriever)
38
+
39
+ def make_inference(query):
40
+ inference = shakespeare_qa.run(query)
41
+ return inference
42
+
43
+ if __name__ == "__main__":
44
+ # make a gradio interface
45
+ import gradio as gr
46
+
47
+ gr.Interface(
48
+ make_inference,
49
+ gr.inputs.Textbox(lines=2, label="Query"),
50
+ gr.outputs.Textbox(label="Response"),
51
+ title="Ask_Shakespeare",
52
+ description="️building_w_llms_qa_Shakespeare allows you to inquire about the Shakespeare's plays.",
53
+ ).launch()