arborvitae commited on
Commit
95fdd56
1 Parent(s): 2f108d3

Upload 4files

Browse files
Files changed (4) hide show
  1. app.py +181 -0
  2. constants.py +8 -0
  3. ingest.py +32 -0
  4. requirements.txt +20 -0
app.py ADDED
@@ -0,0 +1,181 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import streamlit as st
2
+ import os
3
+ import base64
4
+ import time
5
+ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
6
+ from transformers import pipeline
7
+ import torch
8
+ import textwrap
9
+ from langchain.document_loaders import PyPDFLoader, DirectoryLoader, PDFMinerLoader
10
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
11
+ from langchain.embeddings import SentenceTransformerEmbeddings
12
+ from langchain.vectorstores import Chroma
13
+ from langchain.llms import HuggingFacePipeline
14
+ from langchain.chains import RetrievalQA
15
+ from constants import CHROMA_SETTINGS
16
+ from streamlit_chat import message
17
+
18
+ st.set_page_config(layout="wide")
19
+
20
+ device = torch.device('cpu')
21
+
22
+ checkpoint = "MBZUAI/LaMini-T5-738M"
23
+ print(f"Checkpoint path: {checkpoint}") # Add this line for debugging
24
+ tokenizer = AutoTokenizer.from_pretrained(checkpoint)
25
+ base_model = AutoModelForSeq2SeqLM.from_pretrained(
26
+ checkpoint,
27
+ device_map=device,
28
+ torch_dtype=torch.float32
29
+ )
30
+
31
+
32
+ # checkpoint = "LaMini-T5-738M"
33
+ # tokenizer = AutoTokenizer.from_pretrained(checkpoint)
34
+ # base_model = AutoModelForSeq2SeqLM.from_pretrained(
35
+ # checkpoint,
36
+ # device_map="auto",
37
+ # torch_dtype = torch.float32,
38
+ # from_tf=True
39
+ # )
40
+
41
+ persist_directory = "db"
42
+
43
+ @st.cache_resource
44
+ def data_ingestion():
45
+ for root, dirs, files in os.walk("docs"):
46
+ for file in files:
47
+ if file.endswith(".pdf"):
48
+ print(file)
49
+ loader = PDFMinerLoader(os.path.join(root, file))
50
+ documents = loader.load()
51
+ text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=500)
52
+ texts = text_splitter.split_documents(documents)
53
+ #create embeddings here
54
+ embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
55
+ #create vector store here
56
+ db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
57
+ db.persist()
58
+ db=None
59
+
60
+ @st.cache_resource
61
+ def llm_pipeline():
62
+ pipe = pipeline(
63
+ 'text2text-generation',
64
+ model = base_model,
65
+ tokenizer = tokenizer,
66
+ max_length = 256,
67
+ do_sample = True,
68
+ temperature = 0.3,
69
+ top_p= 0.95,
70
+ device=device
71
+ )
72
+ local_llm = HuggingFacePipeline(pipeline=pipe)
73
+ return local_llm
74
+
75
+ @st.cache_resource
76
+ def qa_llm():
77
+ llm = llm_pipeline()
78
+ embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
79
+ db = Chroma(persist_directory="db", embedding_function = embeddings, client_settings=CHROMA_SETTINGS)
80
+ retriever = db.as_retriever()
81
+ qa = RetrievalQA.from_chain_type(
82
+ llm = llm,
83
+ chain_type = "stuff",
84
+ retriever = retriever,
85
+ return_source_documents=True
86
+ )
87
+ return qa
88
+
89
+ def process_answer(instruction):
90
+ response = ''
91
+ instruction = instruction
92
+ qa = qa_llm()
93
+ generated_text = qa(instruction)
94
+ answer = generated_text['result']
95
+ return answer
96
+
97
+ def get_file_size(file):
98
+ file.seek(0, os.SEEK_END)
99
+ file_size = file.tell()
100
+ file.seek(0)
101
+ return file_size
102
+
103
+ @st.cache_data
104
+ #function to display the PDF of a given file
105
+ def displayPDF(file):
106
+ # Opening file from file path
107
+ with open(file, "rb") as f:
108
+ base64_pdf = base64.b64encode(f.read()).decode('utf-8')
109
+
110
+ # Embedding PDF in HTML
111
+ pdf_display = F'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
112
+
113
+ # Displaying File
114
+ st.markdown(pdf_display, unsafe_allow_html=True)
115
+
116
+ # Display conversation history using Streamlit messages
117
+ def display_conversation(history):
118
+ for i in range(len(history["generated"])):
119
+ message(history["past"][i], is_user=True, key=str(i) + "_user")
120
+ message(history["generated"][i],key=str(i))
121
+
122
+ def main():
123
+ st.markdown("<h1 style='text-align: center; color: blue;'>Chat with your PDF 🦜📄 </h1>", unsafe_allow_html=True)
124
+ st.markdown("<h3 style='text-align: center; color: grey;'>Built by <a href='https://github.com/AIAnytime'>AI Anytime with ❤️ </a></h3>", unsafe_allow_html=True)
125
+
126
+ st.markdown("<h2 style='text-align: center; color:red;'>Upload your PDF 👇</h2>", unsafe_allow_html=True)
127
+
128
+ uploaded_file = st.file_uploader("", type=["pdf"])
129
+
130
+ if uploaded_file is not None:
131
+ file_details = {
132
+ "Filename": uploaded_file.name,
133
+ "File size": get_file_size(uploaded_file)
134
+ }
135
+ filepath = "docs/"+uploaded_file.name
136
+ with open(filepath, "wb") as temp_file:
137
+ temp_file.write(uploaded_file.read())
138
+
139
+ col1, col2= st.columns([1,2])
140
+ with col1:
141
+ st.markdown("<h4 style color:black;'>File details</h4>", unsafe_allow_html=True)
142
+ st.json(file_details)
143
+ st.markdown("<h4 style color:black;'>File preview</h4>", unsafe_allow_html=True)
144
+ pdf_view = displayPDF(filepath)
145
+
146
+ with col2:
147
+ with st.spinner('Embeddings are in process...'):
148
+ ingested_data = data_ingestion()
149
+ st.success('Embeddings are created successfully!')
150
+ st.markdown("<h4 style color:black;'>Chat Here</h4>", unsafe_allow_html=True)
151
+
152
+
153
+ user_input = st.text_input("", key="input")
154
+
155
+ # Initialize session state for generated responses and past messages
156
+ if "generated" not in st.session_state:
157
+ st.session_state["generated"] = ["I am ready to help you"]
158
+ if "past" not in st.session_state:
159
+ st.session_state["past"] = ["Hey there!"]
160
+
161
+ # Search the database for a response based on user input and update session state
162
+ if user_input:
163
+ answer = process_answer({'query': user_input})
164
+ st.session_state["past"].append(user_input)
165
+ response = answer
166
+ st.session_state["generated"].append(response)
167
+
168
+ # Display conversation history using Streamlit messages
169
+ if st.session_state["generated"]:
170
+ display_conversation(st.session_state)
171
+
172
+
173
+
174
+
175
+
176
+
177
+
178
+ if __name__ == "__main__":
179
+ main()
180
+
181
+
constants.py ADDED
@@ -0,0 +1,8 @@
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import chromadb
3
+ from chromadb.config import Settings
4
+ CHROMA_SETTINGS = Settings(
5
+ chroma_db_impl='duckdb+parquet',
6
+ persist_directory='db',
7
+ anonymized_telemetry=False
8
+ )
ingest.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from langchain.document_loaders import PyPDFLoader, DirectoryLoader, PDFMinerLoader
2
+ from langchain.text_splitter import RecursiveCharacterTextSplitter
3
+ from langchain.embeddings import SentenceTransformerEmbeddings
4
+ from langchain.vectorstores import Chroma
5
+ import os
6
+ from constants import CHROMA_SETTINGS
7
+
8
+ persist_directory = "db"
9
+
10
+ def main():
11
+ for root, dirs, files in os.walk("docs"):
12
+ for file in files:
13
+ if file.endswith(".pdf"):
14
+ print(file)
15
+ loader = PyPDFLoader(os.path.join(root, file))
16
+ documents = loader.load()
17
+ print("splitting into chunks")
18
+ text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
19
+ texts = text_splitter.split_documents(documents)
20
+ #create embeddings here
21
+ print("Loading sentence transformers model")
22
+ embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
23
+ #create vector store here
24
+ print(f"Creating embeddings. May take some minutes...")
25
+ db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
26
+ db.persist()
27
+ db=None
28
+
29
+ print(f"Ingestion complete! You can now run privateGPT.py to query your documents")
30
+
31
+ if __name__ == "__main__":
32
+ main()
requirements.txt ADDED
@@ -0,0 +1,20 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ langchain==0.0.267
2
+ streamlit==1.25.0
3
+ transformers==4.31.0
4
+ torch==2.0.1
5
+ einops==0.6.1
6
+ bitsandbytes==0.41.1
7
+ accelerate==0.21.0
8
+ pdfminer.six==20221105
9
+ bs4==0.0.1
10
+ sentence_transformers
11
+ duckdb==0.7.1
12
+ chromadb==0.3.26
13
+ beautifulsoup4==4.12.2
14
+ sentence-transformers==2.2.2
15
+ sentencepiece==0.1.99
16
+ six==1.16.0
17
+ requests==2.31.0
18
+ uvicorn==0.18.3
19
+ torch==2.0.1
20
+ torchvision==0.15.2