Add searching for a specific word

app.py

@@ -83,6 +83,7 @@ DATA_DICT, DATA_DESCRIPTION = load_hf_dataset()
 header = st.container()
 description = st.container()
 dataframe_head = st.container()
+word_searching = st.container()
 dataset_statistics = st.container()
 
 with header:
@@ -96,11 +97,17 @@ with dataframe_head:
     filtering_options = DATA_DICT["train"]["target"].unique().tolist()
     filtering_options.append("All classes")
 
-    st.header("First 10 observations of train subset")
+    st.header("First 10 observations of a chosen class")
     class_to_show = st.selectbox(
         label="Select class to show", options=filtering_options
     )
-    df_to_show = DATA_DICT["train"].copy()
+    df_to_show = pd.concat(
+        [
+            DATA_DICT["train"].copy(),
+            DATA_DICT["validation"].copy(),
+            DATA_DICT["test"].copy(),
+        ]
+    )
     if class_to_show == "All classes":
         df_to_show = df_to_show.head(10)
     else:
@@ -108,6 +115,20 @@ with dataframe_head:
     st.dataframe(df_to_show)
     st.text_area(label="Latex code", value=df_to_show.style.to_latex())
 
+with word_searching:
+    st.header("Observations containing a chosen word")
+    searched_word = st.text_input(label="Enter the word you are looking for below")
+    df_to_show = pd.concat(
+        [
+            DATA_DICT["train"].copy(),
+            DATA_DICT["validation"].copy(),
+            DATA_DICT["test"].copy(),
+        ]
+    )
+    df_to_show = df_to_show.loc[df_to_show["text"].str.contains(searched_word)]
+    st.dataframe(df_to_show)
+    st.text_area(label="Latex code", value=df_to_show.style.to_latex())
+
 with dataset_statistics:
     st.header("Dataset statistics")
     st.subheader("Number of samples in each data split")