app.py

"""SecEmbed dataset builder Space.

Builds training pairs, the CyberSec retrieval benchmark, and the demo index
on Hugging Face CPU infrastructure, then pushes all artifacts to the Hub.
Use Rebuild to re-run after code updates.
"""

import io
import threading
import traceback
from contextlib import redirect_stderr, redirect_stdout

import gradio as gr

LOG = io.StringIO()
STATE = {"status": "idle"}
_lock = threading.Lock()


def _run_build() -> None:
    if not _lock.acquire(blocking=False):
        return
    try:
        STATE["status"] = "running"
        LOG.seek(0)
        LOG.truncate(0)
        with redirect_stdout(LOG), redirect_stderr(LOG):
            import importlib
            import build_dataset

            importlib.reload(build_dataset)
            build_dataset.main()
        STATE["status"] = "done"
    except Exception:
        LOG.write("\n" + traceback.format_exc())
        STATE["status"] = "failed"
    finally:
        _lock.release()


def start_build() -> tuple[str, str]:
    if STATE["status"] == "running":
        return STATE["status"], LOG.getvalue()[-24000:]
    threading.Thread(target=_run_build, daemon=True).start()
    return "starting", LOG.getvalue()[-24000:]


def poll() -> tuple[str, str]:
    return STATE["status"], LOG.getvalue()[-24000:]


# Auto-start on boot
threading.Thread(target=_run_build, daemon=True).start()

with gr.Blocks(title="SecEmbed Dataset Builder") as demo:
    gr.Markdown(
        "# SecEmbed Dataset Builder\n"
        "Builds `secembed-pairs`, `cybersec-retrieval-benchmark`, and "
        "`secembed-retrieval-index`, then publishes them to the Hub."
    )
    status = gr.Textbox(label="Status")
    log = gr.Textbox(label="Build log", lines=32)
    with gr.Row():
        refresh = gr.Button("Refresh")
        rebuild = gr.Button("Rebuild", variant="primary")
    refresh.click(poll, outputs=[status, log], api_name="poll")
    rebuild.click(start_build, outputs=[status, log], api_name="rebuild")
    demo.load(poll, outputs=[status, log])

demo.launch()