Nikhil0987 commited on
Commit
49ea6ed
·
verified ·
1 Parent(s): 0fac4e5

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +0 -170
app.py CHANGED
@@ -1,170 +0,0 @@
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
- # from pydantic_settings import BaseSettings
18
-
19
- st.set_page_config(layout="wide")
20
-
21
- device = torch.device('cpu')
22
-
23
- checkpoint = "MBZUAI/LaMini-T5-738M"
24
- print(f"Checkpoint path: {checkpoint}") # Add this line for debugging
25
- tokenizer = AutoTokenizer.from_pretrained(checkpoint)
26
- base_model = AutoModelForSeq2SeqLM.from_pretrained(
27
- checkpoint,
28
- device_map=device,
29
- torch_dtype=torch.float32
30
- )
31
-
32
- persist_directory = "db"
33
-
34
- @st.cache_resource
35
- def data_ingestion():
36
- for root, dirs, files in os.walk("docs"):
37
- for file in files:
38
- if file.endswith(".pdf"):
39
- print(file)
40
- loader = PDFMinerLoader(os.path.join(root, file))
41
- documents = loader.load()
42
- text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=500)
43
- texts = text_splitter.split_documents(documents)
44
- #create embeddings here
45
- embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
46
- #create vector store here
47
- db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
48
- db.persist()
49
- db=None
50
-
51
- @st.cache_resource
52
- def llm_pipeline():
53
- pipe = pipeline(
54
- 'text2text-generation',
55
- model = base_model,
56
- tokenizer = tokenizer,
57
- max_length = 256,
58
- do_sample = True,
59
- temperature = 0.3,
60
- top_p= 0.95,
61
- device=device
62
- )
63
- local_llm = HuggingFacePipeline(pipeline=pipe)
64
- return local_llm
65
-
66
- @st.cache_resource
67
- def qa_llm():
68
- llm = llm_pipeline()
69
- embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
70
- db = Chroma(persist_directory="db", embedding_function = embeddings, client_settings=CHROMA_SETTINGS)
71
- retriever = db.as_retriever()
72
- qa = RetrievalQA.from_chain_type(
73
- llm = llm,
74
- chain_type = "stuff",
75
- retriever = retriever,
76
- return_source_documents=True
77
- )
78
- return qa
79
-
80
- def process_answer(instruction):
81
- response = ''
82
- instruction = instruction
83
- qa = qa_llm()
84
- generated_text = qa(instruction)
85
- answer = generated_text['result']
86
- return answer
87
-
88
- def get_file_size(file):
89
- file.seek(0, os.SEEK_END)
90
- file_size = file.tell()
91
- file.seek(0)
92
- return file_size
93
-
94
- @st.cache_data
95
- #function to display the PDF of a given file
96
- def displayPDF(file):
97
- # Opening file from file path
98
- with open(file, "rb") as f:
99
- base64_pdf = base64.b64encode(f.read()).decode('utf-8')
100
-
101
- # Embedding PDF in HTML
102
- pdf_display = F'<iframe src="data:application/pdf;base64,{base64_pdf}" width="100%" height="600" type="application/pdf"></iframe>'
103
-
104
- # Displaying File
105
- st.markdown(pdf_display, unsafe_allow_html=True)
106
-
107
- # Display conversation history using Streamlit messages
108
- def display_conversation(history):
109
- for i in range(len(history["generated"])):
110
- message(history["past"][i], is_user=True, key=str(i) + "_user")
111
- message(history["generated"][i],key=str(i))
112
-
113
- def main():
114
- st.markdown("<h1 style='text-align: center; color: blue;'>Chat with your PDF 🦜📄 </h1>", unsafe_allow_html=True)
115
- 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)
116
-
117
- st.markdown("<h2 style='text-align: center; color:red;'>Upload your PDF 👇</h2>", unsafe_allow_html=True)
118
-
119
- uploaded_file = st.file_uploader("", type=["pdf"])
120
-
121
- if uploaded_file is not None:
122
- file_details = {
123
- "Filename": uploaded_file.name,
124
- "File size": get_file_size(uploaded_file)
125
- }
126
- filepath = "docs/"+uploaded_file.name
127
- with open(filepath, "wb") as temp_file:
128
- temp_file.write(uploaded_file.read())
129
-
130
- col1, col2= st.columns([1,2])
131
- with col1:
132
- st.markdown("<h4 style color:black;'>File details</h4>", unsafe_allow_html=True)
133
- st.json(file_details)
134
- st.markdown("<h4 style color:black;'>File preview</h4>", unsafe_allow_html=True)
135
- pdf_view = displayPDF(filepath)
136
-
137
- with col2:
138
- with st.spinner('Embeddings are in process...'):
139
- ingested_data = data_ingestion()
140
- st.success('Embeddings are created successfully!')
141
- st.markdown("<h4 style color:black;'>Chat Here</h4>", unsafe_allow_html=True)
142
-
143
-
144
- user_input = st.text_input("", key="input")
145
-
146
- # Initialize session state for generated responses and past messages
147
- if "generated" not in st.session_state:
148
- st.session_state["generated"] = ["I am ready to help you"]
149
- if "past" not in st.session_state:
150
- st.session_state["past"] = ["Hey there!"]
151
-
152
- # Search the database for a response based on user input and update session state
153
- if user_input:
154
- answer = process_answer({'query': user_input})
155
- st.session_state["past"].append(user_input)
156
- response = answer
157
- st.session_state["generated"].append(response)
158
-
159
- # Display conversation history using Streamlit messages
160
- if st.session_state["generated"]:
161
- display_conversation(st.session_state)
162
-
163
-
164
-
165
-
166
-
167
-
168
-
169
- if __name__ == "__main__":
170
- main()