Update app.py

app.py

@@ -9,54 +9,46 @@ 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
 
-def get_all_fields(data):
-    all_keys = set()
-    for record in data:
-        all_keys.update(record.keys())
-    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 ---
+# --- Load session data ---
 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 = []
+
+    # 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.all_fields = get_all_fields(st.session_state.data)
+    st.session_state.prev_data = st.session_state.data.copy()
 
 # --- Upload JSONL File ---
 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 = get_all_fields(st.session_state.data)
+    st.session_state.all_fields = sorted(set().union(*(record.keys() for record in 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.success(
-        f"Loaded {len(st.session_state.data)} records with fields: {st.session_state.all_fields}"
-    )
+    st.success(f"Loaded {len(st.session_state.data)} records.")
 
-# Fallback default fields
+# --- Safe fallback fields if no data yet ---
 if not st.session_state.data and not st.session_state.all_fields:
     st.session_state.all_fields = ["context", "question", "answer"]
 
@@ -67,41 +59,40 @@ df = pd.DataFrame(st.session_state.data)
 df = df.reindex(columns=st.session_state.all_fields)
 
 for field in st.session_state.all_fields:
-    if field.lower() in ["context", "answer", "question"]:
+    if field.lower() in ["context", "question", "answer"]:
         df[field] = df[field].astype(str)
 
 column_configs = {
     field: (
         st.column_config.TextColumn(label=field, width="large")
-        if field.lower() in ["context", "answer", "question"]
+        if field.lower() in ["context", "question", "answer"]
         else None
     )
     for field in st.session_state.all_fields
 }
 
-edited_df = st.data_editor(
-    df,
-    use_container_width=True,
-    num_rows="dynamic",
-    column_config=column_configs,
-    key="editable_table"
-)
+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",
+    )
 
-# --- Auto-save if any changes ---
-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()
+    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")
+        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.toast("✅ Changes auto-saved!", icon="💾")
+        st.toast("✅ Changes auto-saved!", icon="💾")
 
 # --- Add New Entry ---
 st.markdown("### ➕ Add New Entry")
-
 with st.form("new_entry_form"):
     new_record = {}
     for field in st.session_state.all_fields:
@@ -110,6 +101,7 @@ 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:
@@ -127,46 +119,47 @@ with st.expander("�� Add New Field"):
             st.success(f"✅ Field '{new_field}' added!")
             st.rerun()
 
-# --- Export JSONL ---
+# --- 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",
 )
 
-col1, col2 = st.columns(2)
+col1, col2, col3 = st.columns(3)
 
+# Export to path and download
 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_out:
+        with open(export_path, "w", encoding="utf-8") as f:
             for row in st.session_state.data:
-                f_out.write(json.dumps(row, ensure_ascii=False) + "\n")
+                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_download:
-            exported_content = f_download.read()
+        with open(export_path, "r", encoding="utf-8") as f:
+            content = f.read()
 
         st.download_button(
             "⬇️ Download JSONL",
-            exported_content,
+            content,
             file_name=os.path.basename(export_path),
             mime="application/json",
         )
 
+        # Clear session + 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.success("🧹 Temporary session cleared.")
         st.rerun()
 
+# Temp file download
 with col2:
     if os.path.exists(TMP_FILE):
-        with open(TMP_FILE, "r", encoding="utf-8") as f_tmp:
-            tmp_content = f_tmp.read()
-
+        with open(TMP_FILE, "r", encoding="utf-8") as f:
+            tmp_content = f.read()
         st.download_button(
             "⬇️ Download Temp File",
             tmp_content,
@@ -175,3 +168,12 @@ with col2:
         )
     else:
         st.warning("⚠️ No temp file found to download.")
+
+# Clear session button
+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()