Update app.py

app.py

@@ -8,6 +8,7 @@ st.title("📚 JSONL Dataset Editor")
 
 TMP_DIR = "temp"
 TMP_FILE = os.path.join(TMP_DIR, "session_dataset.jsonl")
+META_FILE = os.path.join(TMP_DIR, "metadata.json")
 
 # --- Helper: ensure tmp dir exists ---
 os.makedirs(TMP_DIR, exist_ok=True)
@@ -21,163 +22,203 @@ def get_all_fields(data):
     return sorted(all_keys)
 
 
+if st.button("🔄 Reset Session"):
+    st.session_state.clear()  # Clear session state
+    if os.path.exists(TMP_FILE):
+        os.remove(TMP_FILE)  # Remove temporary file
+    if os.path.exists(META_FILE):
+        os.remove(META_FILE)  # Remove metadata file
+    st.success("🧹 Session has been reset. Starting fresh!")
+    st.rerun()  # Rerun the app to reset everything
+
 # --- Load session data from temp file if exists ---
 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)
     else:
         st.session_state.data = []
-        st.session_state.all_fields = []
+
+    if os.path.exists(META_FILE):
+        with open(META_FILE, "r", encoding="utf-8") as f:
+            metadata = json.load(f)
+            st.session_state.all_fields = [
+                f["name"] for f in metadata.get("fields", [])
+            ]
+            st.session_state.field_types = {
+                f["name"]: f.get("type", "text") for f in metadata.get("fields", [])
+            }
+    else:
+        st.session_state.all_fields = get_all_fields(st.session_state.data)
+        st.session_state.field_types = {
+            field: "text" for field in st.session_state.all_fields
+        }
 
 # --- Upload JSONL File ---
 uploaded_file = st.file_uploader("Upload a JSONL file", type=["jsonl"])
 
 if uploaded_file:
+    # Read the uploaded content
     content = uploaded_file.read().decode("utf-8")
     st.session_state.data = [json.loads(line) for line in content.strip().splitlines()]
+
+    # Update all fields and field types
     st.session_state.all_fields = get_all_fields(st.session_state.data)
 
-    # Save to temp
+    # Update field_types based on metadata if it exists, otherwise set them to 'text'
+    if "field_types" in st.session_state:
+        for field in st.session_state.all_fields:
+            if field not in st.session_state.field_types:
+                st.session_state.field_types[field] = (
+                    "textarea"  # default type for new fields
+                )
+    else:
+        # Initialize field_types for all fields with "text" if no metadata exists
+        st.session_state.field_types = {
+            field: "text" for field in st.session_state.all_fields
+        }
+
+    # Save the updated data and metadata
     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 metadata (field types)
+    metadata = {
+        "fields": [
+            {"name": field, "type": st.session_state.field_types.get(field, "text")}
+            for field in st.session_state.all_fields
+        ]
+    }
+    with open(META_FILE, "w", encoding="utf-8") as f:
+        json.dump(metadata, f, indent=2, ensure_ascii=False)
+
+    # Success message
     st.success(
         f"Loaded {len(st.session_state.data)} records with fields: {st.session_state.all_fields}"
     )
 
 # 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"]
+    st.session_state.all_fields = []
 
 # --- Edit Existing Records ---
-st.markdown("### ✏️ Edit Records")
+if st.session_state.data:
+    st.markdown("### ✏️ Edit Records")
 
-df = pd.DataFrame(st.session_state.data)
-df = df.reindex(columns=st.session_state.all_fields)
+    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:
-    if field.lower() in ["context", "answer", "question"]:
+    # Convert all columns to string for safe editing
+    for field in st.session_state.all_fields:
         df[field] = df[field].astype(str)
 
-# Auto-set long fields like "context", "answer" as textareas
-column_configs = {
-    field: (
-        st.column_config.TextColumn(label=field, width="large")
-        if field.lower() in ["context", "answer", "question"]
-        else None
+    # Ensure field_types is initialized
+    if "field_types" not in st.session_state:
+        st.session_state.field_types = {
+            field: "text" for field in st.session_state.all_fields
+        }
+
+    # Build dynamic column config based on metadata
+    column_configs = {
+        field: (
+            st.column_config.TextColumn(label=field, width="large")
+            if st.session_state.field_types.get(field) == "textarea"
+            else st.column_config.TextColumn(label=field)
+        )
+        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,
     )
-    for field in st.session_state.all_fields
-}
-
-# --- Use st.data_editor for editable table ---
-edited_df = st.data_editor(
-    df,
-    use_container_width=True,
-    num_rows="dynamic",
-    column_config=column_configs,
-)
-
-# 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 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
+    if not edited_df.equals(df):
+        st.session_state.data = edited_df.fillna("").to_dict(orient="records")
         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.toast("✅ Auto-saved!", icon="💾")
         st.rerun()
 
-# Option to add a new field
+# --- Add New Entry ---
+if st.session_state.all_fields:
+    st.markdown("### ➕ Add New Entry")
+
+    with st.form("new_entry_form"):
+        new_record = {}
+
+        # Clear form fields if reset flag is set
+        if "reset_form" in st.session_state and st.session_state.reset_form:
+            # Clear the session state for each input field
+            for field in st.session_state.all_fields:
+                st.session_state[f"input_{field}"] = ""
+            st.session_state.reset_form = (
+                False  # Reset the reset flag after clearing fields
+            )
+
+        # Collect new record input from the user
+        for field in st.session_state.all_fields:
+            input_type = st.session_state.field_types.get(field, "text")
+            if input_type == "textarea":
+                new_record[field] = st.text_area(f"{field}", key=f"input_{field}")
+            else:
+                new_record[field] = st.text_input(f"{field}", key=f"input_{field}")
+
+        submitted = st.form_submit_button("Add Entry")
+
+        if submitted:
+            # Append the new record to session data
+            st.session_state.data.append(new_record)
+
+            # Save the updated data to a 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")
+
+            # Set the reset flag to clear the form next time
+            st.session_state.reset_form = True
+
+            # Show success message and re-run to clear form fields
+            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")
+    new_field = st.text_input("Field name", key="new_field_name")
+    new_type = st.selectbox("Field type", ["text", "textarea"], key="new_field_type")
     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.session_state.field_types[new_field] = new_type
+
+            fields_metadata = [
+                {"name": f, "type": st.session_state.field_types[f]}
+                for f in st.session_state.all_fields
+            ]
+            with open(META_FILE, "w", encoding="utf-8") as f:
+                json.dump({"fields": fields_metadata}, f, indent=2, ensure_ascii=False)
             st.success(f"✅ Field '{new_field}' added!")
             st.rerun()
 
+# --- Download Dataset Button ---
+st.markdown("### 📤 Download Dataset")
 
-# --- 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",
+# Read the session data as a JSONL string
+dataset_content = "\n".join(
+    [json.dumps(row, ensure_ascii=False) for row in st.session_state.data]
 )
 
-col1, col2 = st.columns(2)
-
-# --- Export Button ---
-with col1:
-    if st.button("📁 Export JSONL"):
-        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_out.write(json.dumps(row, ensure_ascii=False) + "\n")
-
-        st.success(f"✅ Dataset saved to {export_path}")
-
-        # 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.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.")
+if os.path.exists(TMP_FILE):
+    # Provide the download button
+    st.download_button(
+        label="⬇️ Download Dataset as JSONL",
+        data=dataset_content,
+        file_name="session_dataset.jsonl",
+        mime="application/json",
+    )
+else:
+    st.warning("Dataset not yet generated!")