Update pages/2_QA.py

pages/2_QA.py

@@ -1,6 +1,4 @@
 import streamlit as st
-import pandas as pd
-import matplotlib.pyplot as plt
 from transformers import pipeline
 
 @st.cache_resource
@@ -12,38 +10,47 @@ def load_qa_model():
 
 qa_model = load_qa_model()
 
+st.title("Question Answering")
+
+if "df" not in st.session_state:
+    st.warning("Please upload and classify a CSV on the Classification page first.")
+    st.stop()
+
+df = st.session_state["df"]
+
+if "content" not in df.columns:
+    st.error("The dataframe must contain a `content` column.")
+    st.stop()
+
 st.markdown("---")
-            st.subheader("Question Answering")
-
-            article_labels = [
-                f"Article {i + 1}: {text[:100]}{'...' if len(text) > 100 else ''}"
-                for i, text in enumerate(df["content"])
-            ]
-
-            selected_label = st.selectbox("Choose an article", article_labels)
-            selected_index = article_labels.index(selected_label)
-            selected_article = df["content"].iloc[selected_index]
-
-            st.write("### Selected Article")
-            st.write(selected_article)
-
-            question = st.text_input("Ask a question about the selected article")
-
-            if st.button("Get Answer"):
-                if not question.strip():
-                    st.warning("Please enter a question.")
-                elif not selected_article.strip():
-                    st.warning("The selected article is empty.")
-                else:
-                    with st.spinner("Finding answer..."):
-                        result = qa_model(
-                            question=question,
-                            context=selected_article[:3000]
-                        )
-
-                    st.success("Answer found.")
-                    st.write("### Answer")
-                    st.write(result["answer"])
-
-    except Exception as e:
-        st.error(f"Error reading or processing file: {e}")
+st.subheader("Question Answering")
+
+article_labels = [
+    f"Article {i + 1}: {text[:100]}{'...' if len(text) > 100 else ''}"
+    for i, text in enumerate(df["content"].fillna("").astype(str))
+]
+
+selected_label = st.selectbox("Choose an article", article_labels)
+selected_index = article_labels.index(selected_label)
+selected_article = df["content"].fillna("").astype(str).iloc[selected_index]
+
+st.write("### Selected Article")
+st.write(selected_article)
+
+question = st.text_input("Ask a question about the selected article")
+
+if st.button("Get Answer"):
+    if not question.strip():
+        st.warning("Please enter a question.")
+    elif not selected_article.strip():
+        st.warning("The selected article is empty.")
+    else:
+        with st.spinner("Finding answer..."):
+            result = qa_model(
+                question=question,
+                context=selected_article[:3000]
+            )
+
+        st.success("Answer found.")
+        st.write("### Answer")
+        st.write(result["answer"])