Update app.py

app.py

@@ -12,11 +12,7 @@ from streamlit_chat import message
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 import torch
 
-# Constants
-CHECKPOINT = "MBZUAI/LaMini-T5-738M"
-TOKENIZER = AutoTokenizer.from_pretrained(CHECKPOINT)
-BASE_MODEL = AutoModelForSeq2SeqLM.from_pretrained(CHECKPOINT, device_map=torch.device('cpu'), torch_dtype=torch.float32)
-
+st.set_page_config(layout="wide")
 
 def process_answer(instruction, qa_chain):
     response = ''
@@ -50,7 +46,11 @@ def data_ingestion():
 
 
 @st.cache_resource
-def initialize_qa_chain():
+def initialize_qa_chain(selected_model):
+    # Constants
+    CHECKPOINT = selected_model
+    TOKENIZER = AutoTokenizer.from_pretrained(CHECKPOINT)
+    BASE_MODEL = AutoModelForSeq2SeqLM.from_pretrained(CHECKPOINT, device_map=torch.device('cpu'), torch_dtype=torch.float32)
     pipe = pipeline(
         'text2text-generation',
         model=BASE_MODEL,
@@ -101,7 +101,10 @@ def display_conversation(history):
 
 
 def main():
-    st.set_page_config(layout="wide")
+     # Add a sidebar for model selection
+    model_options = ["MBZUAI/LaMini-T5-738M", "google/flan-t5-base", "google/flan-t5-small"]
+    selected_model = st.sidebar.selectbox("Select Model", model_options)
+    
     st.markdown("<h1 style='text-align: center; color: blue;'>Custom PDF Chatbot 🦜📄 </h1>", unsafe_allow_html=True)
     st.markdown("<h2 style='text-align: center; color:red;'>Upload your PDF, and Ask Questions 👇</h2>", unsafe_allow_html=True)
 
@@ -125,6 +128,7 @@ def main():
                 pdf_view = display_pdf(filepath)
 
             with col2:
+                st.success(f'model selected successfully: {selected_model}')
                 with st.spinner('Embeddings are in process...'):
                     ingested_data = data_ingestion()
                 st.success('Embeddings are created successfully!')
@@ -140,7 +144,7 @@ def main():
 
                 # Search the database for a response based on user input and update session state
                 if user_input:
-                    answer = process_answer({'query': user_input}, initialize_qa_chain())
+                    answer = process_answer({'query': user_input}, initialize_qa_chain(selected_model))
                     st.session_state["past"].append(user_input)
                     response = answer
                     st.session_state["generated"].append(response)