More improvements

app.py

@@ -2,10 +2,8 @@ from datetime import datetime
 
 import streamlit as st
 import pandas as pd
-import numpy as np
 import matplotlib.pyplot as plt
 
-from datasets import Dataset
 from load_dataframe import get_data
 
 
@@ -49,8 +47,8 @@ def aggregated_data(df, aggregation_level="week"):
     st.pyplot(plt)
 
 
-def show_data_editor(df: pd.DataFrame, key: str):
-    edited_df = st.data_editor(df,
+def show_data_editor(filtered_df: pd.DataFrame, key: str):
+    edited_df = st.data_editor(filtered_df,
                 hide_index=True,
                 column_order=("reached_out", "reached_out_link", "paper_page", "title", "github", "num_models", "num_datasets", "num_spaces"),
                 column_config={"github": st.column_config.LinkColumn(),
@@ -59,47 +57,38 @@ def show_data_editor(df: pd.DataFrame, key: str):
                 width=2000,
                 key=key)
     
-    # Check if the dataframe has been edited
-    # TODO this is wrong
-    # rather we should probably do a merge-join (overwriting the edited rows) and then save the new dataframe
-    # if not edited_df.equals(df):
-    #     save_data(edited_df)
-    #     st.success("Changes saved successfully!")
+    if edited_df is not None and not edited_df.equals(filtered_df):
+        # update the df of the session state with the affected rows
+        # TODO there seems to be a bug in here
+        original_df = st.session_state.df
+        original_df.update(edited_df)
+        st.session_state.df = original_df
 
 
-def save_data(df: pd.DataFrame):
-    # load as HF dataset
-    dataset = Dataset.from_pandas(df)
-
-    raise NotImplementedError("To do")
-
-
-def display_data(df: pd.DataFrame):
-    df['has_artifact'] = (df['num_models'] > 0) | (df['num_datasets'] > 0) | (df['num_spaces'] > 0)
-    num_artifacts = df['has_artifact'].sum()
-    percentage_of_at_least_one_artifact = num_artifacts / df.shape[0] if df.shape[0] > 0 else 0
+def display_data(filtered_df: pd.DataFrame):
+    num_artifacts = filtered_df['has_artifact'].sum()
+    percentage_of_at_least_one_artifact = num_artifacts / filtered_df.shape[0] if filtered_df.shape[0] > 0 else 0
     percentage_of_at_least_one_artifact = round(percentage_of_at_least_one_artifact * 100, 2)
 
-    # add reached out and reached out link columns
-    df['reached_out'] = [False for _ in range(df.shape[0])]
-    df["reached_out_link"] = ["" for _ in range(df.shape[0])]
-
     st.markdown(f"""
     ## {percentage_of_at_least_one_artifact}% papers with at least one 🤗 artifact
     
-    * Number of papers: {df.shape[0]}
-    * Number of papers with a Github link: {df['github'].notnull().sum()}
+    * Number of papers: {filtered_df.shape[0]}
+    * Number of papers with a Github link: {(filtered_df['github'].values != '').sum()}
     * Number of papers with at least one HF artifact: {num_artifacts}
     """)
 
     st.write("Papers with at least one artifact")
-    show_data_editor(df[df['has_artifact']], key="papers_with_artifacts")
+    show_data_editor(filtered_df=filtered_df[filtered_df['has_artifact']],
+                     key="papers_with_artifacts")
 
     st.write("Papers without artifacts")
-    show_data_editor(df[~df['has_artifact']], key="papers_without_artifacts")
+    show_data_editor(filtered_df=filtered_df[~filtered_df['has_artifact']],
+                     key="papers_without_artifacts")
     
     st.write("Papers with a HF mention in README but no artifacts")
-    show_data_editor(df[(df['hf_mention'] == 1) & (~df['has_artifact'])], key="papers_with_hf_mention_no_artifacts")
+    show_data_editor(filtered_df=filtered_df[(filtered_df['hf_mention'] == 1) & (~filtered_df['has_artifact'])],
+                     key="papers_with_hf_mention_no_artifacts")
 
 
 def main():
@@ -109,31 +98,41 @@ def main():
     st.sidebar.title("Navigation")
     selection = st.sidebar.selectbox("Go to", ["Daily/weekly/monthly data", "Aggregated data"])
 
+    # Initialize session state
+    if 'df' not in st.session_state:
+        df = get_data()
+        # add has_artifact, reached out and reached out link columns
+        # TODO remove since this will overwrite everything if we have added data before
+        df['has_artifact'] = (df['num_models'] > 0) | (df['num_datasets'] > 0) | (df['num_spaces'] > 0)
+        df['reached_out'] = [False for _ in range(df.shape[0])]
+        df["reached_out_link"] = ["" for _ in range(df.shape[0])]
+        st.session_state.df = df
+
     if selection == "Daily/weekly/monthly data":
         # Button to select day, month or week
         # Add streamlit selectbox.
         view_level = st.selectbox(label="View data per day, week or month", options=["day", "week", "month"])
 
         if view_level == "day":
-            # get the latest dataframe
-            df = get_data()
-            
             # make a button to select the day, defaulting to today
             day = st.date_input("Select day", value="today", format="DD/MM/YYYY")
             # convert to the day of a Pandas Timestamp
             day = pd.Timestamp(day)
 
+            # fetch df from sessions state
+            df = st.session_state.df
+
             filtered_df = df[df.index.date == day.date()]
 
             st.write(f"Showing data for {day.day_name()} {day.strftime('%d/%m/%Y')}")
-            display_data(df=filtered_df)
+            display_data(filtered_df=filtered_df)
 
-        elif view_level == "week":
-            # get the latest dataframe
-            df = get_data()
-            
+        elif view_level == "week":            
             # make a button to select the week
             week_number = st.number_input("Select week", value=datetime.today().isocalendar()[1], min_value=1, max_value=52)
+
+            # fetch df from sessions state
+            df = st.session_state.df
             
             # Extract week number from the index
             df['week'] = df.index.isocalendar().week
@@ -143,15 +142,15 @@ def main():
             
             st.write(f"Showing data for week {week_number}")
             
-            display_data(df=filtered_df)
+            display_data(filtered_df=filtered_df)
 
-        elif view_level == "month":
-            # get the latest dataframe
-            df = get_data()
-            
+        elif view_level == "month":            
             # make a button to select the month, defaulting to current month
             month_str = st.selectbox("Select month", options=["January", "February", "March", "April", "May", "June", "July", "August", "September", "October", "November", "December"])
             year_str = st.selectbox("Select year", options=["2024"])
+
+            # fetch df from sessions state
+            df = st.session_state.df
             
             # Filter the dataframe for the desired week number
             month_map = {
@@ -167,7 +166,7 @@ def main():
             
             st.write(f"Showing data for {month_str} {year_str}")
             
-            display_data(df=filtered_df)
+            display_data(filtered_df=filtered_df)
 
     elif selection == "Aggregated data":