Update app.py

app.py

@@ -66,7 +66,7 @@ with st.sidebar:
 st.title("Check Your Paper Now!")
 
 # Main content
-option = st.radio("Select input type:",("Text", "PDF"))
+option = st.radio("Select input type:",("Text"))
 
 if option == "Text":
     title_input = st.text_area("Enter Title:")
@@ -82,20 +82,31 @@ if option == "Text":
         selected_category = custom_category if custom_category else "Other"
 
     combined_text = f"{title_input} [SEP] {keywords_input} [SEP] {abstract_input} [SEP] {selected_category} [SEP] {affiliations_input}"
-    if st.button("Predict"):
-        if not any([title_input, abstract_input,keywords_input,  affiliations_input]):
+        if st.button("Predict"):
+        if not any([title_input, abstract_input, keywords_input, full_text_input, affiliations_input]):
             st.warning("Please enter paper text.")
         else:
             with st.spinner("Predicting..."):
-                predicted_class = predict_class(combined_text)
-                if predicted_class is not None:
+                longformer_class = predict_class(combined_text, longformer_model, longformer_tokenizer)
+                bert_class = predict_class(combined_text, bert_model, bert_tokenizer)
+                if longformer_class is not None and bert_class is not None:
                     class_labels = ["Level 1", "Level 2", "Level 3", "Level 4"]
 
-                    st.text("Predicted Class:")
+                    st.text("Longformer Predicted Class:")
                     for i, label in enumerate(class_labels):
-                        if i == predicted_class:
+                        if i == longformer_class:
                             st.markdown(
-                                f'<div style="background-color: {class_colors[predicted_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
+                                f'<div style="background-color: {class_colors[longformer_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
+                                unsafe_allow_html=True
+                            )
+                        else:
+                            st.text(label)
+
+                    st.text("BERT Predicted Class:")
+                    for i, label in enumerate(class_labels):
+                        if i == bert_class:
+                            st.markdown(
+                                f'<div style="background-color: {class_colors[bert_class]}; padding: 10px; border-radius: 5px; color: white; font-weight: bold;">{label}</div>',
                                 unsafe_allow_html=True
                             )
                         else: