Update app.py

app.py

@@ -9,11 +9,9 @@ st.title("📚 JSONL Dataset Editor")
 TMP_DIR = "temp"
 TMP_FILE = os.path.join(TMP_DIR, "session_dataset.jsonl")
 
-# --- Helper: ensure tmp dir exists ---
 os.makedirs(TMP_DIR, exist_ok=True)
 
 
-# --- Helper: get all unique fields from records ---
 def get_all_fields(data):
     all_keys = set()
     for record in data:
@@ -21,15 +19,25 @@ def get_all_fields(data):
     return sorted(all_keys)
 
 
+# --- Clear session handler ---
+if st.button("🧹 Clear Session"):
+    st.session_state.clear()
+    if os.path.exists(TMP_FILE):
+        os.remove(TMP_FILE)
+    st.success("Session and temp file cleared!")
+    st.rerun()
+
 # --- Load session data from temp file if exists ---
 if "data" not in st.session_state:
     if os.path.exists(TMP_FILE):
         with open(TMP_FILE, "r", encoding="utf-8") as f:
             st.session_state.data = [json.loads(line) for line in f]
         st.session_state.all_fields = get_all_fields(st.session_state.data)
+        st.session_state.prev_data = st.session_state.data.copy()
     else:
         st.session_state.data = []
         st.session_state.all_fields = []
+        st.session_state.prev_data = []
 
 # --- Upload JSONL File ---
 uploaded_file = st.file_uploader("Upload a JSONL file", type=["jsonl"])
@@ -38,8 +46,8 @@ if uploaded_file:
     content = uploaded_file.read().decode("utf-8")
     st.session_state.data = [json.loads(line) for line in content.strip().splitlines()]
     st.session_state.all_fields = get_all_fields(st.session_state.data)
+    st.session_state.prev_data = st.session_state.data.copy()
 
-    # Save to temp
     with open(TMP_FILE, "w", encoding="utf-8") as f:
         for item in st.session_state.data:
             f.write(json.dumps(item, ensure_ascii=False) + "\n")
@@ -48,7 +56,7 @@ if uploaded_file:
         f"Loaded {len(st.session_state.data)} records with fields: {st.session_state.all_fields}"
     )
 
-# If still no data, use safe fallback fields
+# Fallback default fields
 if not st.session_state.data and not st.session_state.all_fields:
     st.session_state.all_fields = ["context", "question", "answer"]
 
@@ -58,12 +66,10 @@ st.markdown("### ✏️ Edit Records")
 df = pd.DataFrame(st.session_state.data)
 df = df.reindex(columns=st.session_state.all_fields)
 
-# Fix: Convert likely text fields to string to avoid StreamlitAPIException
 for field in st.session_state.all_fields:
     if field.lower() in ["context", "answer", "question"]:
         df[field] = df[field].astype(str)
 
-# Auto-set long fields like "context", "answer" as textareas
 column_configs = {
     field: (
         st.column_config.TextColumn(label=field, width="large")
@@ -73,49 +79,38 @@ column_configs = {
     for field in st.session_state.all_fields
 }
 
-# --- Use st.data_editor for editable table ---
-# Track previous data
-if "prev_data" not in st.session_state:
-    st.session_state.prev_data = []
-
-# Show editable table
 edited_df = st.data_editor(
     df,
     use_container_width=True,
     num_rows="dynamic",
     column_config=column_configs,
-    key="editable_table",
+    key="editable_table"
 )
 
-# Save only if changed
+# --- Auto-save if any changes ---
 new_data = edited_df.fillna("").to_dict(orient="records")
-
-if new_data != st.session_state.get("prev_data", []):
+if new_data != st.session_state.prev_data:
     st.session_state.data = new_data
-    st.session_state.prev_data = new_data
+    st.session_state.prev_data = new_data.copy()
 
     with open(TMP_FILE, "w", encoding="utf-8") as f:
-        for item in new_data:
+        for item in st.session_state.data:
             f.write(json.dumps(item, ensure_ascii=False) + "\n")
 
-    st.experimental_rerun()
-
-
+    st.toast("✅ Changes auto-saved!", icon="💾")
 
 # --- Add New Entry ---
 st.markdown("### ➕ Add New Entry")
 
-# Show form with current fields
 with st.form("new_entry_form"):
     new_record = {}
     for field in st.session_state.all_fields:
         new_record[field] = st.text_area(f"{field}", key=f"input_{field}")
-
     submitted = st.form_submit_button("Add Entry")
     if submitted:
         st.session_state.data.append(new_record)
+        st.session_state.prev_data = st.session_state.data.copy()
 
-        # Save to temp
         with open(TMP_FILE, "w", encoding="utf-8") as f:
             for item in st.session_state.data:
                 f.write(json.dumps(item, ensure_ascii=False) + "\n")
@@ -123,7 +118,7 @@ with st.form("new_entry_form"):
         st.success("✅ New entry added!")
         st.rerun()
 
-# Option to add a new field
+# --- Add New Field ---
 with st.expander("➕ Add New Field"):
     new_field = st.text_input("New field name", key="new_field_name")
     if st.button("Add Field"):
@@ -132,11 +127,9 @@ with st.expander("➕ Add New Field"):
             st.success(f"✅ Field '{new_field}' added!")
             st.rerun()
 
-
 # --- Export JSONL ---
 st.markdown("### 📤 Export Dataset")
 
-# Let user define a custom export path
 export_path = st.text_input(
     "Custom save path (e.g., ./exports/my_dataset.jsonl)",
     value="./exports/exported_dataset.jsonl",
@@ -144,20 +137,15 @@ export_path = st.text_input(
 
 col1, col2 = st.columns(2)
 
-# --- Export Button ---
 with col1:
     if st.button("📁 Export JSONL"):
-        if not os.path.exists(os.path.dirname(export_path)):
-            os.makedirs(os.path.dirname(export_path))
-
-        # Write to custom path
+        os.makedirs(os.path.dirname(export_path), exist_ok=True)
         with open(export_path, "w", encoding="utf-8") as f_out:
             for row in st.session_state.data:
                 f_out.write(json.dumps(row, ensure_ascii=False) + "\n")
 
         st.success(f"✅ Dataset saved to {export_path}")
 
-        # Load content for download
         with open(export_path, "r", encoding="utf-8") as f_download:
             exported_content = f_download.read()
 
@@ -168,14 +156,12 @@ with col1:
             mime="application/json",
         )
 
-        # Reset session and temp
         if os.path.exists(TMP_FILE):
             os.remove(TMP_FILE)
         st.session_state.clear()
         st.success("🧹 Temporary session cleared. You're starting fresh!")
         st.rerun()
 
-# --- Download Temp Only Button ---
 with col2:
     if os.path.exists(TMP_FILE):
         with open(TMP_FILE, "r", encoding="utf-8") as f_tmp: