Update app.py

app.py

@@ -2,151 +2,182 @@ 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")
 
 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)
 
 
-# --- Helpers ---
+# --- Helper: get all unique fields from records ---
 def get_all_fields(data):
-    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 row in st.session_state.data:
-            f.write(json.dumps(row, ensure_ascii=False) + "\n")
+    all_keys = set()
+    for record in data:
+        all_keys.update(record.keys())
+    return sorted(all_keys)
 
 
-# --- Session Initialization ---
+# --- Load session data from temp file if exists ---
 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 "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 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)
-
-# --- 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]
+    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)
+    else:
+        st.session_state.data = []
+        st.session_state.all_fields = []
+
+# --- 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)
-    save_to_file()
-    st.session_state.editor_key = str(uuid4())
-    st.session_state.just_added = True
-    st.rerun()  # Force rerun to update editor
 
-# --- Add New Fields fallback ---
-if not st.session_state.all_fields:
-    st.session_state.all_fields = ["context", "question", "answer"]
-
-# --- 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")
+    # 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")
 
-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()  # Force rerun to update editor
+    st.success(
+        f"Loaded {len(st.session_state.data)} records with fields: {st.session_state.all_fields}"
+    )
 
-# --- Wait for rerun before rendering editor ---
-if st.session_state.just_added:
-    st.session_state.just_added = False
-    st.rerun()
+# If still no data, use safe fallback fields
+if not st.session_state.data and not st.session_state.all_fields:
+    st.session_state.all_fields = ["context", "question", "answer"]
 
-# --- Display Editor ---
+# --- Edit Existing Records ---
 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:
-    df[field] = df[field].astype(str)
+    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(field, width="large")
-    if field.lower() in ["context", "question", "answer"]
-    else None
+    field: (
+        st.column_config.TextColumn(label=field, width="large")
+        if field.lower() in ["context", "answer", "question"]
+        else None
+    )
     for field in st.session_state.all_fields
 }
 
+# --- Use st.data_editor for editable table ---
 edited_df = st.data_editor(
     df,
-    key=st.session_state.editor_key,
     use_container_width=True,
     num_rows="dynamic",
     column_config=column_configs,
 )
 
-# --- Save edits immediately after editing ---
-if edited_df is not None:
-    new_data = edited_df.fillna("").to_dict(orient="records")
-    if new_data != st.session_state.data:  # Ensure the data has changed before saving
-        st.session_state.data = new_data
-        save_to_file()
-        st.toast("✅ Changes auto-saved!", icon="💾")
+# Auto-save logic: detect changes and persist
+if not edited_df.equals(df):
+    st.session_state.data = edited_df.fillna("").to_dict(orient="records")
+
+    # Save to temp 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")
+
+    st.toast("✅ Auto-saved!", icon="💾")
+    st.rerun()
+
 
-# --- Add New Field ---
+# --- 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)
+
+        # 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")
+
+        st.success("✅ New entry added!")
+        st.rerun()
+
+# Option to add a 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"):
         if new_field and new_field not in st.session_state.all_fields:
             st.session_state.all_fields.append(new_field)
-            st.rerun()  # Rerun after adding field
+            st.success(f"✅ Field '{new_field}' added!")
+            st.rerun()
 
-# --- Export Dataset ---
-st.markdown("### 📤 Export")
-export_path = st.text_input("Save path", value="./exports/exported_dataset.jsonl")
 
-col1, col2, col3 = st.columns(3)
+# --- 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",
+)
+
+col1, col2 = st.columns(2)
+
+# --- Export Button ---
 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:
+        if not os.path.exists(os.path.dirname(export_path)):
+            os.makedirs(os.path.dirname(export_path))
+
+        # Write to custom path
+        with open(export_path, "w", encoding="utf-8") as f_out:
             for row in st.session_state.data:
-                f.write(json.dumps(row, ensure_ascii=False) + "\n")
-        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.rerun()  # Force rerun after export
+                f_out.write(json.dumps(row, ensure_ascii=False) + "\n")
 
-with col2:
-    if os.path.exists(TMP_FILE):
-        with open(TMP_FILE, "r", encoding="utf-8") as f:
-            tmp_data = f.read()
-        st.download_button("⬇️ Temp File", tmp_data, file_name="session_dataset.jsonl")
+        st.success(f"✅ Dataset saved to {export_path}")
 
-with col3:
-    if st.button("🧹 Clear Session"):
+        # Load content for download
+        with open(export_path, "r", encoding="utf-8") as f_download:
+            exported_content = f_download.read()
+
+        st.download_button(
+            "⬇️ Download JSONL",
+            exported_content,
+            file_name=os.path.basename(export_path),
+            mime="application/json",
+        )
+
+        # Reset session and temp
         if os.path.exists(TMP_FILE):
             os.remove(TMP_FILE)
         st.session_state.clear()
-        st.rerun()  # Rerun after clearing session
+        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:
+            tmp_content = f_tmp.read()
+
+        st.download_button(
+            "⬇️ Download Temp File",
+            tmp_content,
+            file_name="session_dataset.jsonl",
+            mime="application/json",
+        )
+    else:
+        st.warning("⚠️ No temp file found to download.")