Update app.py

app.py

@@ -1,19 +1,15 @@
-import streamlit as st 
+import streamlit as st
 import os
 import base64
 import time
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM 
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import pipeline
-import torch 
-import textwrap 
-from langchain.document_loaders import PyPDFLoader, DirectoryLoader, PDFMinerLoader 
-from langchain.text_splitter import RecursiveCharacterTextSplitter 
-from langchain.embeddings import SentenceTransformerEmbeddings 
-from langchain.vectorstores import Chroma 
+import torch
+import textwrap
+from langchain.embeddings import SentenceTransformerEmbeddings
+from langchain.vectorstores import Chroma
 from langchain.llms import HuggingFacePipeline
-from langchain.chains import RetrievalQA 
-from constants import CHROMA_SETTINGS
-from streamlit_chat import message
+from langchain.chains import RetrievalQA
 
 st.set_page_config(layout="wide")
 
@@ -28,45 +24,31 @@ base_model = AutoModelForSeq2SeqLM.from_pretrained(
     torch_dtype=torch.float32
 )
 
-
-# checkpoint = "LaMini-T5-738M"
-# tokenizer = AutoTokenizer.from_pretrained(checkpoint)
-# base_model = AutoModelForSeq2SeqLM.from_pretrained(
-#     checkpoint,
-#     device_map="auto",
-#     torch_dtype = torch.float32,
-#     from_tf=True
-# )
-
 persist_directory = "db"
 
 @st.cache_resource
-def data_ingestion():
-    for root, dirs, files in os.walk("docs"):
-        for file in files:
-            if file.endswith(".pdf"):
-                print(file)
-                loader = PDFMinerLoader(os.path.join(root, file))
-    documents = loader.load()
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=500)
-    texts = text_splitter.split_documents(documents)
-    #create embeddings here
-    embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
-    #create vector store here
-    db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
-    db.persist()
-    db=None 
+def data_ingestion(uploaded_file):
+    with open(uploaded_file.name, "rb") as pdf_file:
+        pdf_content = pdf_file.read()
+
+    # Process the PDF content here and generate a brief summary.
+    # You can use libraries like PyPDF2, pdfminer, or other PDF processing tools.
+
+    # For now, let's assume we have extracted the text from the PDF.
+    pdf_text = "This is a brief summary of the PDF content."
+
+    return pdf_text
 
 @st.cache_resource
 def llm_pipeline():
     pipe = pipeline(
         'text2text-generation',
-        model = base_model,
-        tokenizer = tokenizer,
-        max_length = 256,
-        do_sample = True,
-        temperature = 0.3,
-        top_p= 0.95,
+        model=base_model,
+        tokenizer=tokenizer,
+        max_length=256,
+        do_sample=True,
+        temperature=0.3,
+        top_p=0.95,
         device=device
     )
     local_llm = HuggingFacePipeline(pipeline=pipe)
@@ -76,12 +58,12 @@ def llm_pipeline():
 def qa_llm():
     llm = llm_pipeline()
     embeddings = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
-    db = Chroma(persist_directory="db", embedding_function = embeddings, client_settings=CHROMA_SETTINGS)
+    db = Chroma(persist_directory="db", embedding_function=embeddings, client_settings=CHROMA_SETTINGS)
     retriever = db.as_retriever()
     qa = RetrievalQA.from_chain_type(
-        llm = llm,
-        chain_type = "stuff",
-        retriever = retriever,
+        llm=llm,
+        chain_type="stuff",
+        retriever=retriever,
         return_source_documents=True
     )
     return qa
@@ -101,7 +83,7 @@ def get_file_size(file):
     return file_size
 
 @st.cache_data
-#function to display the PDF of a given file 
+# Function to display the PDF of a given file
 def displayPDF(file):
     # Opening file from file path
     with open(file, "rb") as f:
@@ -117,7 +99,7 @@ def displayPDF(file):
 def display_conversation(history):
     for i in range(len(history["generated"])):
         message(history["past"][i], is_user=True, key=str(i) + "_user")
-        message(history["generated"][i],key=str(i))
+        message(history["generated"][i], key=str(i))
 
 def main():
     st.markdown("<h1 style='text-align: center; color: blue;'>Chat with your PDF 🦜📄 </h1>", unsafe_allow_html=True)
@@ -132,49 +114,20 @@ def main():
             "Filename": uploaded_file.name,
             "File size": get_file_size(uploaded_file)
         }
-        filepath = "docs/"+uploaded_file.name
-        with open(filepath, "wb") as temp_file:
-                temp_file.write(uploaded_file.read())
 
-        col1, col2= st.columns([1,2])
+        col1, col2 = st.columns([1, 2])
         with col1:
             st.markdown("<h4 style color:black;'>File details</h4>", unsafe_allow_html=True)
             st.json(file_details)
             st.markdown("<h4 style color:black;'>File preview</h4>", unsafe_allow_html=True)
-            pdf_view = displayPDF(filepath)
+            pdf_view = displayPDF(uploaded_file)
 
         with col2:
-            with st.spinner('Embeddings are in process...'):
-                ingested_data = data_ingestion()
-            st.success('Embeddings are created successfully!')
-            st.markdown("<h4 style color:black;'>Chat Here</h4>", unsafe_allow_html=True)
-
-
-            user_input = st.text_input("", key="input")
-
-            # Initialize session state for generated responses and past messages
-            if "generated" not in st.session_state:
-                st.session_state["generated"] = ["I am ready to help you"]
-            if "past" not in st.session_state:
-                st.session_state["past"] = ["Hey there!"]
-                
-            # Search the database for a response based on user input and update session state
-            if user_input:
-                answer = process_answer({'query': user_input})
-                st.session_state["past"].append(user_input)
-                response = answer
-                st.session_state["generated"].append(response)
-
-            # Display conversation history using Streamlit messages
-            if st.session_state["generated"]:
-                display_conversation(st.session_state)
-        
-
-        
-
-
-
+            with st.spinner('Processing the uploaded PDF...'):
+                pdf_summary = data_ingestion(uploaded_file)
+            st.success('PDF processing is complete!')
+            st.markdown("<h4 style color:black;'>PDF Summary</h4>", unsafe_allow_html=True)
+            st.write(pdf_summary)
 
 if __name__ == "__main__":
     main()
-