from __future__ import annotations from typing import Any import pandas as pd import streamlit as st from data_studio.auth import authenticate_user, count_users, get_active_user from data_studio.config import ConfigError, StudioSettings from data_studio.db import create_client, ensure_indexes, get_database from data_studio.importer import import_source_samples from data_studio.review import apply_review_action, claim_next_sample, queue_status_counts from data_studio.sources import list_sources, seed_default_sources from data_studio.utils import normalize_text, split_csv_tags, split_document_lines @st.cache_resource def load_settings() -> StudioSettings: return StudioSettings.from_env() @st.cache_resource def load_database() -> Any: settings = load_settings() client = create_client(settings) return get_database(client, settings) def get_session_user() -> dict[str, Any] | None: username = st.session_state.get("auth_username") if not username: return None return { "username": username, "display_name": st.session_state.get("auth_display_name", username), "role": st.session_state.get("auth_role", "reviewer"), } def set_session_user(user: dict[str, Any]) -> None: st.session_state["auth_username"] = user["username"] st.session_state["auth_display_name"] = user.get("display_name") or user["username"] st.session_state["auth_role"] = user.get("role", "reviewer") def clear_session_user() -> None: for key in ("auth_username", "auth_display_name", "auth_role"): st.session_state.pop(key, None) def render_login(db) -> None: st.title("Rachana Data Studio") st.caption("Invite-only access for approved reviewers.") if count_users(db) == 0: st.warning( "No users are configured yet. Create the first invite-only account with " "`python .\\scripts\\create_data_studio_user.py --username `." ) with st.form("login-form", clear_on_submit=False): username = st.text_input("Username") password = st.text_input("Password", type="password") submitted = st.form_submit_button("Log in", type="primary", use_container_width=True) if submitted: user = authenticate_user(db, username, password) if user is None: st.error("Invalid username or password, or the account is inactive.") else: set_session_user(user) st.rerun() def require_authenticated_user(db) -> dict[str, Any] | None: session_user = get_session_user() if session_user is None: return None active_user = get_active_user(db, session_user["username"]) if active_user is None: clear_session_user() st.warning("Your account is no longer active. Please log in again.") st.rerun() return { "username": active_user["username"], "display_name": active_user.get("display_name") or active_user["username"], "role": active_user.get("role", "reviewer"), } def render_overview(db) -> None: counts = queue_status_counts(db) st.subheader("Queue overview") if counts: total_pending = sum(statuses.get("pending", 0) for statuses in counts.values()) total_accepted = sum(statuses.get("accepted", 0) for statuses in counts.values()) total_edited = sum(statuses.get("edited", 0) for statuses in counts.values()) total_rejected = sum(statuses.get("rejected", 0) for statuses in counts.values()) cards = st.columns(4) cards[0].metric("Pending", total_pending) cards[1].metric("Accepted", total_accepted) cards[2].metric("Edited", total_edited) cards[3].metric("Rejected", total_rejected) rows: list[dict[str, Any]] = [] for queue_name, statuses in counts.items(): row = {"queue_name": queue_name} row.update(statuses) rows.append(row) st.dataframe(pd.DataFrame(rows).fillna(0), width="stretch") else: st.info("No samples imported yet.") st.subheader("Source configs") source_rows = list_sources(db) if source_rows: st.dataframe( pd.DataFrame( [ { "dataset_key": row["dataset_key"], "queue_name": row["queue_name"], "sample_type": row["sample_type"], "hf_dataset": row["hf_dataset"], "hf_config": row.get("hf_config"), "hf_split": row["hf_split"], } for row in source_rows ] ), width="stretch", ) else: st.warning("No source configs seeded yet.") def render_sources(db) -> None: st.subheader("Source setup") st.caption("Seed the default source configurations for the active review queues.") if st.button("Seed default sources", type="primary"): inserted = seed_default_sources(db) st.success(f"Source config seed completed. New records inserted: {inserted}.") st.rerun() source_rows = list_sources(db) if source_rows: for row in source_rows: with st.container(border=True): top = st.columns([2, 2, 2, 1]) top[0].markdown(f"**{row['name']}**") top[1].caption(f"Queue: `{row['queue_name']}`") top[2].caption(f"Type: `{row['sample_type']}`") top[3].caption("Enabled" if row.get("enabled", True) else "Disabled") st.caption(f"Dataset key: `{row['dataset_key']}`") st.caption(f"HF dataset: `{row['hf_dataset']}`") if row.get("description"): st.write(row["description"]) def render_import(db, settings: StudioSettings) -> None: st.subheader("Import Hugging Face source data") sources = list_sources(db) if not sources: st.warning("Seed source configs first.") return dataset_keys = [source["dataset_key"] for source in sources if source.get("enabled", True)] selected = st.selectbox("Source", dataset_keys) selected_source = next(source for source in sources if source["dataset_key"] == selected) current_counts = queue_status_counts(db).get(selected_source["queue_name"], {}) info_cols = st.columns(4) info_cols[0].metric("Queue", selected_source["queue_name"]) info_cols[1].metric("Pending now", current_counts.get("pending", 0)) info_cols[2].metric("Accepted", current_counts.get("accepted", 0)) info_cols[3].metric("Edited", current_counts.get("edited", 0)) st.caption(f"HF dataset: `{selected_source['hf_dataset']}`") max_samples = st.number_input("Max samples to import", min_value=10, max_value=50000, value=1000, step=10) quick_cols = st.columns(4) if quick_cols[0].button("Import 25", use_container_width=True): max_samples = 25 st.session_state["import_now"] = True if quick_cols[1].button("Import 100", use_container_width=True): max_samples = 100 st.session_state["import_now"] = True if quick_cols[2].button("Import 500", use_container_width=True): max_samples = 500 st.session_state["import_now"] = True manual_import = quick_cols[3].button("Import custom", type="primary", use_container_width=True) if manual_import or st.session_state.pop("import_now", False): with st.spinner("Importing samples from Hugging Face..."): stats = import_source_samples(db, settings, selected, int(max_samples)) st.success("Import complete.") st.json(stats) def render_document_editor(sample: dict[str, Any]) -> dict[str, Any]: sample_key = sample["sample_id"] with st.expander("Original document", expanded=False): st.text_area("Original", value=sample["payload"]["text"], height=220, disabled=True) current_lines = split_document_lines(sample["review"]["current_text"]) line_df = pd.DataFrame( { "keep": [True] * len(current_lines), "line_no": list(range(1, len(current_lines) + 1)), "text": current_lines, } ) edited_df = st.data_editor( line_df, width="stretch", hide_index=True, num_rows="fixed", column_config={ "keep": st.column_config.CheckboxColumn("Keep"), "line_no": st.column_config.NumberColumn("Line", disabled=True), "text": st.column_config.TextColumn("Text", width="large"), }, disabled=["line_no"], key=f"{sample['sample_id']}-line-editor", ) kept_lines = [ normalize_text(str(row["text"])) for _, row in edited_df.iterrows() if bool(row["keep"]) and normalize_text(str(row["text"])) ] added_lines_raw = st.text_area( "Add new lines (one per line)", value="", height=100, key=f"{sample_key}-added-lines", ) added_lines = [normalize_text(line) for line in added_lines_raw.splitlines()] kept_lines.extend(line for line in added_lines if line) current_text = "\n".join(kept_lines) st.caption(f"Kept lines: {len(kept_lines)}") st.text_area("Final cleaned text preview", value=current_text, height=180, disabled=True) return {"current_text": current_text} def render_translation_editor(sample: dict[str, Any]) -> dict[str, Any]: col1, col2 = st.columns(2) with col1: st.markdown("**Original source**") st.text_area("Original source text", value=sample["payload"]["source_text"], height=180, disabled=True) source_text = st.text_area( "Reviewed source text", value=sample["review"]["current_pair"]["source_text"], height=180, ) with col2: st.markdown("**Original target**") st.text_area("Original target text", value=sample["payload"]["target_text"], height=180, disabled=True) target_text = st.text_area( "Reviewed target text", value=sample["review"]["current_pair"]["target_text"], height=180, ) return {"current_pair": {"source_text": source_text, "target_text": target_text}} def render_transliteration_editor(sample: dict[str, Any]) -> dict[str, Any]: col1, col2 = st.columns(2) with col1: st.text_input("Original native text", value=sample["payload"]["native_text"], disabled=True) native_text = st.text_input( "Reviewed native text", value=sample["review"]["current_transliteration"]["native_text"], ) with col2: st.text_input("Original Latin text", value=sample["payload"]["latin_text"], disabled=True) latin_text = st.text_input( "Reviewed Latin text", value=sample["review"]["current_transliteration"]["latin_text"], ) return {"current_transliteration": {"native_text": native_text, "latin_text": latin_text}} def render_review(db, user: dict[str, Any]) -> None: st.subheader("Review queue") counts = queue_status_counts(db) sources = [source for source in list_sources(db) if source.get("enabled", True)] dataset_options = ["All active datasets"] + [source["dataset_key"] for source in sources] selected_dataset = st.selectbox("Dataset", dataset_options) queue_name = "all" queue_stats = counts.get("pure_telugu", {}) stats_cols = st.columns(4) stats_cols[0].metric("Pending", queue_stats.get("pending", 0)) stats_cols[1].metric("Accepted", queue_stats.get("accepted", 0)) stats_cols[2].metric("Edited", queue_stats.get("edited", 0)) stats_cols[3].metric("Rejected", queue_stats.get("rejected", 0)) sample = claim_next_sample( db, queue_name, user["username"], None if selected_dataset == "All active datasets" else selected_dataset, ) if sample is None: st.info("No pending samples for the selected pool.") return st.caption(f"Sample ID: {sample['sample_id']}") meta_cols = st.columns(4) meta_cols[0].markdown(f"**Dataset**\n\n`{sample['source']['dataset_key']}`") meta_cols[1].markdown(f"**Record ID**\n\n`{sample['source'].get('source_record_id')}`") meta_cols[2].markdown(f"**Status**\n\n`{sample['status']}`") meta_cols[3].markdown(f"**Type**\n\n`{sample['sample_type']}`") if sample["source"].get("source_title"): st.caption(f"Title: {sample['source']['source_title']}") with st.form(f"review-{sample['sample_id']}"): if sample["sample_type"] == "document": cleaned_payload = render_document_editor(sample) elif sample["sample_type"] == "translation_pair": cleaned_payload = render_translation_editor(sample) else: cleaned_payload = render_transliteration_editor(sample) with st.expander("Tags and review notes", expanded=True): keyword_text = st.text_input( "Keywords / tags (comma separated)", value=", ".join(sample["review"]["quality_tags"]), ) task_tag_text = st.text_input( "Task tags (comma separated)", value=", ".join(sample["review"]["task_tags"]), ) notes = st.text_area("Reviewer notes", value=sample["review"].get("notes", ""), height=100) reason = st.text_input("Action reason") confidence = st.slider("Reviewer confidence", min_value=1, max_value=5, value=4) col1, col2, col3, col4 = st.columns(4) accept = col1.form_submit_button("Accept", type="primary", use_container_width=True) edit = col2.form_submit_button("Save Edit", use_container_width=True) reject = col3.form_submit_button("Reject", use_container_width=True) skip = col4.form_submit_button("Skip", use_container_width=True) action = None if accept: action = "accept" elif edit: action = "edit" elif reject: action = "reject" elif skip: action = "skip" if action is not None: apply_review_action( db=db, sample_id=sample["sample_id"], action=action, reviewer=user["username"], expected_version=sample["version"], cleaned_payload=cleaned_payload, tags=split_csv_tags(keyword_text), quality_tags=split_csv_tags(keyword_text), task_tags=split_csv_tags(task_tag_text), notes=notes, confidence=confidence, reason=reason, ) st.success(f"Action recorded: {action}") st.rerun() def main() -> None: st.set_page_config(page_title="Rachana Data Studio", layout="wide") try: settings = load_settings() db = load_database() except ConfigError as exc: st.error(str(exc)) st.stop() ensure_indexes(db) user = require_authenticated_user(db) if user is None: render_login(db) return st.title("Rachana Data Studio") st.caption("HF-native corpus artifacts with MongoDB-backed review state.") st.sidebar.header("Workspace") section = st.sidebar.radio("Section", ["Overview", "Sources", "Import", "Review"]) st.sidebar.caption(f"MongoDB DB: `{settings.mongodb_db_name}`") st.sidebar.caption(f"Signed in as: `{user['display_name']}`") st.sidebar.caption(f"Role: `{user['role']}`") if st.sidebar.button("Log out", use_container_width=True): clear_session_user() st.rerun() if section == "Overview": render_overview(db) elif section == "Sources": render_sources(db) elif section == "Import": render_import(db, settings) else: render_review(db, user) if __name__ == "__main__": main()