Update app.py

app.py

@@ -1,35 +1,31 @@
 import streamlit as st
 from streamlit_chat import message
 from langchain.chains import ConversationalRetrievalChain
-from langchain.chains import RetrievalQA
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.llms import  HuggingFacePipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain.vectorstores import FAISS
 from langchain.memory import ConversationBufferMemory
 from langchain_community.document_loaders import PyPDFLoader
-from transformers import pipeline, T5Tokenizer, T5ForConditionalGeneration
+from transformers import T5Tokenizer, T5ForConditionalGeneration
 import torch
 from transformers import pipeline
 import os
 import tempfile
 
-#model = AutoModelForSeq2SeqLM.from_pretrained("LaMini-Flan-T5-783M")
-#tokenizer = AutoTokenizer.from_pretrained("LaMini-Flan-T5-783M", device_map = 'auto')
 
-
-checkpoint = "LaMini-Flan-T5-783M"
+checkpoint = "LaMini-Flan-T5-248M"
 tokenizer = T5Tokenizer.from_pretrained(checkpoint)
-model = T5ForConditionalGeneration.from_pretrained( checkpoint, device_map = 'auto', torch_dtype = torch.float32 )
+
+base_model = T5ForConditionalGeneration.from_pretrained( checkpoint, device_map = 'auto', torch_dtype = torch.float32 )
 
 def llm_pipeline():
     pipe = pipeline(
         'text2text-generation',
-        model = model,
+        model = base_model,
         tokenizer = tokenizer,
         do_sample = True,
         temperature = 0.5,
-        max_length = 250
     )
     local_llm = HuggingFacePipeline(pipeline=pipe)
     return local_llm
@@ -58,7 +54,6 @@ def display_chat_history(chain):
         with st.form(key='my_form', clear_on_submit=True):
             user_input = st.text_input("Question:", placeholder="Ask about your PDF", key='input')
             submit_button = st.form_submit_button(label='Send')
-            
 
         if submit_button and user_input:
             with st.spinner('Generating response...'):
@@ -81,17 +76,13 @@ def create_conversational_chain(vector_store):
 
     chain = ConversationalRetrievalChain.from_llm(llm=llm, chain_type='stuff',
                                                  retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
-                                                 memory=memory) 
-    """
-    chain = RetrievalQA.from_chain_type(llm=llm, chain_type="stuff",
-                                     retriever=vector_store.as_retriever(search_kwargs={"k": 2}),
-                                    memory = memory)
-    return chain """
+                                                 memory=memory)
+    return chain
 
 def main():
     # Initialize session state
     initialize_session_state()
-    st.title("Multi-PDF ChatBot :books:")
+    st.title("Multi-PDF ChatBot using Mistral-7B-Instruct :books:")
     # Initialize Streamlit
     st.sidebar.title("Document Processing")
     uploaded_files = st.sidebar.file_uploader("Upload files", accept_multiple_files=True)