Update app.py

app.py

@@ -8,60 +8,57 @@ st.title("📚 JSONL Dataset Editor")
 
 TMP_DIR = "temp"
 TMP_FILE = os.path.join(TMP_DIR, "session_dataset.jsonl")
-
-# --- Setup temp directory ---
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# --- Reset update flag on rerun ---
-if st.session_state.get("updated"):
-    st.session_state.updated = False
+# --- Helpers ---
+def get_all_fields(data):
+    all_keys = set()
+    for record in data:
+        all_keys.update(record.keys())
+    return sorted(all_keys)
 
-# --- Load session data ---
+# --- Session Initialization ---
 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]
-    else:
-        st.session_state.data = []
-
-    # Load all unique fields
-    def get_all_fields(data):
-        all_keys = set()
-        for record in data:
-            all_keys.update(record.keys())
-        return sorted(all_keys)
-
+    st.session_state.data = []
+if "all_fields" not in st.session_state:
+    st.session_state.all_fields = []
+if "prev_data" not in st.session_state:
+    st.session_state.prev_data = []
+
+# --- Load from temp if needed ---
+if not st.session_state.data and 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()
 
-# --- Upload JSONL File ---
+# --- Upload JSONL ---
 uploaded_file = st.file_uploader("Upload a JSONL file", type=["jsonl"])
 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 = sorted(set().union(*(record.keys() for record in st.session_state.data)))
+    st.session_state.all_fields = get_all_fields(st.session_state.data)
     st.session_state.prev_data = st.session_state.data.copy()
-
     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")
+    st.rerun()
 
-    st.success(f"Loaded {len(st.session_state.data)} records.")
-
-# --- Safe fallback fields if no data yet ---
-if not st.session_state.data and not st.session_state.all_fields:
+# --- Fallback fields if none ---
+if not st.session_state.all_fields:
     st.session_state.all_fields = ["context", "question", "answer"]
 
-# --- Edit Existing Records ---
+# --- Edit Records ---
 st.markdown("### ✏️ Edit Records")
-
 df = pd.DataFrame(st.session_state.data)
 df = df.reindex(columns=st.session_state.all_fields)
 
+# Ensure fields are strings for editor
 for field in st.session_state.all_fields:
     if field.lower() in ["context", "question", "answer"]:
         df[field] = df[field].astype(str)
 
+# TextAreas for longer fields
 column_configs = {
     field: (
         st.column_config.TextColumn(label=field, width="large")
@@ -71,25 +68,23 @@ column_configs = {
     for field in st.session_state.all_fields
 }
 
-if not st.session_state.get("updated"):
-    edited_df = st.data_editor(
-        df,
-        use_container_width=True,
-        num_rows="dynamic",
-        column_config=column_configs,
-        key="editable_table",
-    )
-
-    new_data = edited_df.fillna("").to_dict(orient="records")
-    if new_data != st.session_state.prev_data:
-        st.session_state.data = new_data
-        st.session_state.prev_data = new_data.copy()
-
-        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")
+edited_df = st.data_editor(
+    df,
+    use_container_width=True,
+    num_rows="dynamic",
+    column_config=column_configs,
+    key="editable_table",
+)
 
-        st.toast("✅ Changes auto-saved!", icon="💾")
+# Save if changed
+new_data = edited_df.fillna("").to_dict(orient="records")
+if new_data != st.session_state.prev_data:
+    st.session_state.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:
+            f.write(json.dumps(item, ensure_ascii=False) + "\n")
+    st.toast("✅ Auto-saved!", icon="💾")
 
 # --- Add New Entry ---
 st.markdown("### ➕ Add New Entry")
@@ -101,12 +96,9 @@ with st.form("new_entry_form"):
     if submitted:
         st.session_state.data.append(new_record)
         st.session_state.prev_data = st.session_state.data.copy()
-        st.session_state.updated = True
-
         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")
-
         st.success("✅ New entry added!")
         st.rerun()
 
@@ -116,46 +108,30 @@ with st.expander("➕ Add New Field"):
     if st.button("Add Field"):
         if new_field and new_field not in st.session_state.all_fields:
             st.session_state.all_fields.append(new_field)
-            st.success(f"✅ Field '{new_field}' added!")
             st.rerun()
 
 # --- Export Section ---
 st.markdown("### 📤 Export Dataset")
-export_path = st.text_input(
-    "Custom save path (e.g., ./exports/my_dataset.jsonl)",
-    value="./exports/exported_dataset.jsonl",
-)
+export_path = st.text_input("Save path", value="./exports/exported_dataset.jsonl")
 
 col1, col2, col3 = st.columns(3)
 
-# Export to path and download
+# Export
 with col1:
     if st.button("📁 Export JSONL"):
         os.makedirs(os.path.dirname(export_path), exist_ok=True)
         with open(export_path, "w", encoding="utf-8") as f:
             for row in st.session_state.data:
                 f.write(json.dumps(row, ensure_ascii=False) + "\n")
-
-        st.success(f"✅ Dataset saved to {export_path}")
-
         with open(export_path, "r", encoding="utf-8") as f:
             content = f.read()
-
-        st.download_button(
-            "⬇️ Download JSONL",
-            content,
-            file_name=os.path.basename(export_path),
-            mime="application/json",
-        )
-
-        # Clear session + temp
+        st.download_button("⬇️ Download JSONL", content, file_name=os.path.basename(export_path))
         if os.path.exists(TMP_FILE):
             os.remove(TMP_FILE)
         st.session_state.clear()
-        st.success("🧹 Temporary session cleared.")
         st.rerun()
 
-# Temp file download
+# Download temp
 with col2:
     if os.path.exists(TMP_FILE):
         with open(TMP_FILE, "r", encoding="utf-8") as f:
@@ -167,13 +143,12 @@ with col2:
             mime="application/json",
         )
     else:
-        st.warning("⚠️ No temp file found to download.")
+        st.warning("⚠️ No temp file found.")
 
-# Clear session button
+# Clear session
 with col3:
     if st.button("🗑️ Clear Session"):
         if os.path.exists(TMP_FILE):
             os.remove(TMP_FILE)
         st.session_state.clear()
-        st.success("🧹 Session and temp cleared!")
         st.rerun()