Update app.py

app.py

@@ -1,11 +1,12 @@
 import os
 import streamlit as st
+from streamlit_option_menu import option_menu
 from langchain.embeddings.openai import OpenAIEmbeddings
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.vectorstores import FAISS 
 from transformers import TFAutoModelForQuestionAnswering, AutoTokenizer, pipeline
 
-os.environ["OPENAI_API_KEY"] = "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"
+os.environ["OPENAI_API_KEY"] = "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"
 
 # Read data
 with open("./data/full_context.txt", "r") as file1:
@@ -20,52 +21,81 @@ text_splitter = CharacterTextSplitter(
 )
 texts = text_splitter.split_text(doc)
 
-
 # Download embeddings from OpenAI
 embeddings = OpenAIEmbeddings()
 docsearch = FAISS.from_texts(texts, embeddings)
 
-# Load model
-model_path = "./models/roberta_model"
-
-model = TFAutoModelForQuestionAnswering.from_pretrained(model_path)
-tokenizer = AutoTokenizer.from_pretrained('deepset/roberta-base-squad2')
+# Load roberta model
+model_path0 = "./models/roberta_model"
+model0 = TFAutoModelForQuestionAnswering.from_pretrained(model_path0)
+tokenizer0 = AutoTokenizer.from_pretrained('deepset/roberta-base-squad2')
+# Initialize Transformer pipeline with our own model and tokenizer
+question_answerer0 = pipeline("question-answering", model=model0, tokenizer=tokenizer0)
 
+# Load bert base model
+model_path1 = "./models/bert_finetuned_model"
+model1 = TFAutoModelForQuestionAnswering.from_pretrained(model_path1)
+tokenizer1 = AutoTokenizer.from_pretrained('huggingface-course/bert-finetuned-squad')
 # Initialize Transformer pipeline with our own model and tokenizer
-question_answerer = pipeline("question-answering", model=model, tokenizer=tokenizer)
+question_answerer1 = pipeline("question-answering", model=model1, tokenizer=tokenizer1)
 
-def findHighestScore(question):
-    docs_found = docsearch.similarity_search(question)
-    doc_score = 0.01
-    doc_answer = ''
 
+def QnAfunction(question, QnAmodel):
+    docs_found = docsearch.similarity_search(question)
+    score = 0.5
+    answer = ''
     for doc in docs_found:
-        doc_result = question_answerer(question=question, context = doc.page_content)
-        if doc_result['score'] > doc_score:
-            doc_score = doc_result['score']
-            doc_answer = doc_result['answer']
-    
-    return doc_answer, doc_score
+        doc_result = QnAmodel(question=question, context = doc.page_content)
+        if doc_result['score'] > score:
+            score = doc_result['score']
+            answer = doc_result['answer']
 
-
-def QnAfunction(question):
-    answer1, score1 = findHighestScore(question)
-    if answer1 != '':
-        return answer1, score1
+    if answer != '':
+        return answer, score
         # print("Answer: ", answer1)
         # print("Score: ", score1)
-
     else:
         return "No Answer found. Please ask question related to Bachelor of Computer Science program at Swinburne.", 0
         # print("No Answer found. Please ask question related to Bachelor of Computer Science program at Swinburne.")
 
-
-text = st.text_area("Ask any question about the Bachelor of Computer Science program at Swinburne: ")
-if text:
-    ans, score = QnAfunction(text)
-    if score > 0.5:
-        st.write("Answer: ", ans)
-        st.write("Score: ", score)
-    else:
-        st.write(ans)
-
+# GUI with Streamlit
+st.markdown("""
+    <style> 
+        .big-font {
+            margin: 50px 0 10px 0 !important;
+            font-size:25px !important;
+            font-weight: bold !important;
+        }
+    </style>
+    """, unsafe_allow_html=True)
+
+with st.sidebar:
+    selected = option_menu("Model selection", ["Roberta base squad2", "Bert finetuned squad"], 
+        icons=['box-fill', 'box-fill'], menu_icon="cast", default_index=0)
+
+if selected == "Roberta base squad2":
+    st.markdown('<p class="big-font">QnA for Swinburne\'s Bachelor of Computer Science progrom</p>', unsafe_allow_html=True)
+    st.write("- ", selected)
+    text = st.text_area("Type question (Eg. What is the duration of the Bachelor of Computer Science program?): ")
+    if text:
+        #######
+        ans, score = QnAfunction(text, question_answerer0)
+        if score > 0.5:
+            st.write("Answer: ", ans)
+            st.write("Score: ", score)
+        else:
+            st.write(ans)
+
+
+elif selected == "Bert finetuned squad":
+    st.markdown('<p class="big-font">QnA for Swinburne\'s Bachelor of Computer Science progrom</p>', unsafe_allow_html=True)
+    st.write("- ", selected)
+    text = st.text_area("Type question (Eg. What is the duration of the Bachelor of Computer Science program?): ")
+    if text:
+        # Fed in the question to the model
+        ans, score = QnAfunction(text, question_answerer1)
+        if score > 0.5:
+            st.write("Answer: ", ans)
+            st.write("Score: ", score)
+        else:
+            st.write(ans)