Update app.py

app.py

@@ -2,6 +2,7 @@ import streamlit as st
 import json
 import pandas as pd
 import os
+from uuid import uuid4
 
 st.set_page_config(page_title="Dataset Builder", layout="wide")
 st.title("📚 JSONL Dataset Editor")
@@ -17,20 +18,24 @@ def get_all_fields(data):
         all_keys.update(record.keys())
     return sorted(all_keys)
 
-# --- Session Initialization ---
+def save_to_file():
+    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")
+
+# --- Initialize state ---
 if "data" not in st.session_state:
     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 = []
+if "edit_key" not in st.session_state:
+    st.session_state.edit_key = str(uuid4())
 
-# --- Load from temp if needed ---
-if not st.session_state.data and os.path.exists(TMP_FILE):
+# --- Load from temp if file exists ---
+if os.path.exists(TMP_FILE) and not st.session_state.data:
     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 ---
 uploaded_file = st.file_uploader("Upload a JSONL file", type=["jsonl"])
@@ -38,27 +43,22 @@ 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()
-    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")
+    save_to_file()
+    st.session_state.edit_key = str(uuid4())
     st.rerun()
 
-# --- Fallback fields if none ---
+# --- Ensure default fields if none ---
 if not st.session_state.all_fields:
     st.session_state.all_fields = ["context", "question", "answer"]
 
-# --- Edit Records ---
+# --- Edit Section ---
 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)
+    df[field] = df[field].astype(str)
 
-# TextAreas for longer fields
 column_configs = {
     field: (
         st.column_config.TextColumn(label=field, width="large")
@@ -70,35 +70,31 @@ column_configs = {
 
 edited_df = st.data_editor(
     df,
+    key=st.session_state.edit_key,  # Unique key forces Streamlit to rerender editor
     use_container_width=True,
     num_rows="dynamic",
     column_config=column_configs,
-    key="editable_table",
 )
 
-# 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="💾")
+# If data changed, update state and save
+if edited_df is not None:
+    new_data = edited_df.fillna("").to_dict(orient="records")
+    if new_data != st.session_state.data:
+        st.session_state.data = new_data
+        save_to_file()
+        st.toast("✅ Changes auto-saved!", icon="💾")
 
-# --- Add New Entry ---
+# --- Add Entry ---
 st.markdown("### ➕ Add New Entry")
-with st.form("new_entry_form"):
-    new_record = {}
+with st.form("add_form"):
+    new_item = {}
     for field in st.session_state.all_fields:
-        new_record[field] = st.text_area(f"{field}", key=f"input_{field}")
+        new_item[field] = st.text_area(field, key=f"add_{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()
-        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.session_state.data.append(new_item)
+        save_to_file()
+        st.session_state.edit_key = str(uuid4())  # Force editor to refresh
         st.success("✅ New entry added!")
         st.rerun()
 
@@ -116,7 +112,6 @@ export_path = st.text_input("Save path", value="./exports/exported_dataset.jsonl
 
 col1, col2, col3 = st.columns(3)
 
-# Export
 with col1:
     if st.button("📁 Export JSONL"):
         os.makedirs(os.path.dirname(export_path), exist_ok=True)
@@ -131,7 +126,6 @@ with col1:
         st.session_state.clear()
         st.rerun()
 
-# Download temp
 with col2:
     if os.path.exists(TMP_FILE):
         with open(TMP_FILE, "r", encoding="utf-8") as f:
@@ -142,10 +136,7 @@ with col2:
             file_name="session_dataset.jsonl",
             mime="application/json",
         )
-    else:
-        st.warning("⚠️ No temp file found.")
 
-# Clear session
 with col3:
     if st.button("🗑️ Clear Session"):
         if os.path.exists(TMP_FILE):