Update app.py

app.py

@@ -11,47 +11,73 @@ TMP_DIR = "temp"
 TMP_FILE = os.path.join(TMP_DIR, "session_dataset.jsonl")
 os.makedirs(TMP_DIR, exist_ok=True)
 
+
 # --- Helpers ---
 def get_all_fields(data):
-    all_keys = set()
-    for record in data:
-        all_keys.update(record.keys())
-    return sorted(all_keys)
+    keys = set()
+    for d in data:
+        keys.update(d.keys())
+    return sorted(list(keys))
+
 
 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")
+        for row in st.session_state.data:
+            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+
 
-# --- Initialize state ---
+# --- Session Initialization ---
 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 "edit_key" not in st.session_state:
-    st.session_state.edit_key = str(uuid4())
+if "editor_key" not in st.session_state:
+    st.session_state.editor_key = str(uuid4())
+if "just_added" not in st.session_state:
+    st.session_state.just_added = False
 
-# --- Load from temp if file exists ---
+# --- Load from TMP file ---
 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)
 
-# --- 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()]
+# --- File Upload ---
+uploaded = st.file_uploader("Upload JSONL", type=["jsonl"])
+if uploaded:
+    lines = uploaded.read().decode("utf-8").splitlines()
+    st.session_state.data = [json.loads(l) for l in lines]
     st.session_state.all_fields = get_all_fields(st.session_state.data)
     save_to_file()
-    st.session_state.edit_key = str(uuid4())
+    st.session_state.editor_key = str(uuid4())
     st.rerun()
 
-# --- Ensure default fields if none ---
+# --- Add New Fields fallback ---
 if not st.session_state.all_fields:
     st.session_state.all_fields = ["context", "question", "answer"]
 
-# --- Edit Section ---
+# --- Add Entry Form ---
+with st.form("add_form"):
+    st.markdown("### ➕ Add New Entry")
+    new_entry = {}
+    for field in st.session_state.all_fields:
+        new_entry[field] = st.text_area(field, key=f"add_{field}")
+
+    submit_add = st.form_submit_button("Add Entry")
+
+if submit_add:
+    st.session_state.data.append(new_entry)
+    save_to_file()
+    st.session_state.editor_key = str(uuid4())
+    st.session_state.just_added = True
+    st.rerun()
+
+# --- Wait for rerun before rendering editor ---
+if st.session_state.just_added:
+    st.session_state.just_added = False
+    st.rerun()
+
+# --- Display Editor ---
 st.markdown("### ✏️ Edit Records")
 df = pd.DataFrame(st.session_state.data)
 df = df.reindex(columns=st.session_state.all_fields)
@@ -60,23 +86,21 @@ for field in st.session_state.all_fields:
     df[field] = df[field].astype(str)
 
 column_configs = {
-    field: (
-        st.column_config.TextColumn(label=field, width="large")
-        if field.lower() in ["context", "question", "answer"]
-        else None
-    )
+    field: st.column_config.TextColumn(field, width="large")
+    if field.lower() in ["context", "question", "answer"]
+    else None
     for field in st.session_state.all_fields
 }
 
 edited_df = st.data_editor(
     df,
-    key=st.session_state.edit_key,  # Unique key forces Streamlit to rerender editor
+    key=st.session_state.editor_key,
     use_container_width=True,
     num_rows="dynamic",
     column_config=column_configs,
 )
 
-# If data changed, update state and save
+# --- Save edits ---
 if edited_df is not None:
     new_data = edited_df.fillna("").to_dict(orient="records")
     if new_data != st.session_state.data:
@@ -84,20 +108,6 @@ if edited_df is not None:
         save_to_file()
         st.toast("✅ Changes auto-saved!", icon="💾")
 
-# --- Add Entry ---
-st.markdown("### ➕ Add New Entry")
-with st.form("add_form"):
-    new_item = {}
-    for field in st.session_state.all_fields:
-        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_item)
-        save_to_file()
-        st.session_state.edit_key = str(uuid4())  # Force editor to refresh
-        st.success("✅ New entry added!")
-        st.rerun()
-
 # --- Add New Field ---
 with st.expander("➕ Add New Field"):
     new_field = st.text_input("New field name", key="new_field_name")
@@ -106,8 +116,8 @@ with st.expander("➕ Add New Field"):
             st.session_state.all_fields.append(new_field)
             st.rerun()
 
-# --- Export Section ---
-st.markdown("### 📤 Export Dataset")
+# --- Export Dataset ---
+st.markdown("### ���� Export")
 export_path = st.text_input("Save path", value="./exports/exported_dataset.jsonl")
 
 col1, col2, col3 = st.columns(3)
@@ -121,24 +131,20 @@ with col1:
         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))
+        st.success("✅ Exported!")
+        st.session_state.clear()
         if os.path.exists(TMP_FILE):
             os.remove(TMP_FILE)
-        st.session_state.clear()
         st.rerun()
 
 with col2:
     if os.path.exists(TMP_FILE):
         with open(TMP_FILE, "r", encoding="utf-8") as f:
-            tmp_content = f.read()
-        st.download_button(
-            "⬇️ Download Temp File",
-            tmp_content,
-            file_name="session_dataset.jsonl",
-            mime="application/json",
-        )
+            tmp_data = f.read()
+        st.download_button("⬇️ Temp File", tmp_data, file_name="session_dataset.jsonl")
 
 with col3:
-    if st.button("🗑️ Clear Session"):
+    if st.button("🧹 Clear Session"):
         if os.path.exists(TMP_FILE):
             os.remove(TMP_FILE)
         st.session_state.clear()