Upload folder using huggingface_hub

README.md

@@ -23,6 +23,16 @@ Set these Space variables in **Settings → Variables and secrets**:
 
 The app will download checkpoint from your model repo and load it at runtime.
 
+### Optional MLflow Tracking in Space
+
+You can enable lightweight MLflow event logging for inference + task runs.
+Set these optional variables in **Settings → Variables and secrets**:
+
+- `MLFLOW_TRACKING_URI` (example: `file:/tmp/mlruns` or your remote tracking server URI)
+- `MLFLOW_EXPERIMENT_NAME` (example: `hf-space-sanskrit-d3pm`)
+
+If not set, the Space runs normally without MLflow.
+
 ## Local Dev
 
 ```bash

app.py

@@ -6,6 +6,7 @@ import sys
 import shutil
 import threading
 import uuid
+import time
 from datetime import datetime
 from pathlib import Path
 
@@ -31,6 +32,45 @@ DEFAULT_ANALYSIS_OUT = "analysis_outputs/T4"
 os.makedirs(RESULTS_DIR, exist_ok=True)
 _BG_JOBS = {}
 
+try:
+    import mlflow
+except Exception:
+    mlflow = None
+
+_MLFLOW_READY = False
+
+
+def _setup_mlflow_once():
+    global _MLFLOW_READY
+    if _MLFLOW_READY:
+        return
+    if mlflow is None:
+        return
+    try:
+        tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "file:/tmp/mlruns")
+        experiment = os.environ.get("MLFLOW_EXPERIMENT_NAME", "hf-space-sanskrit-d3pm")
+        mlflow.set_tracking_uri(tracking_uri)
+        mlflow.set_experiment(experiment)
+        _MLFLOW_READY = True
+    except Exception:
+        _MLFLOW_READY = False
+
+
+def _mlflow_event(run_name: str, params: dict | None = None, metrics: dict | None = None, tags: dict | None = None):
+    _setup_mlflow_once()
+    if not _MLFLOW_READY or mlflow is None:
+        return
+    try:
+        with mlflow.start_run(run_name=run_name, nested=False):
+            if tags:
+                mlflow.set_tags({k: str(v) for k, v in tags.items()})
+            if params:
+                mlflow.log_params({k: (v if isinstance(v, (int, float, str, bool)) else str(v)) for k, v in params.items()})
+            if metrics:
+                mlflow.log_metrics({k: float(v) for k, v in metrics.items()})
+    except Exception:
+        pass
+
 HF_DEFAULT_MODEL_REPO = os.environ.get("HF_DEFAULT_MODEL_REPO", "bhsinghgrid/DevaFlow")
 HF_DEFAULT_MODEL_FILE = os.environ.get("HF_DEFAULT_MODEL_FILE", "best_model.pt")
 
@@ -259,6 +299,7 @@ def generate_from_ui(
     if not input_text.strip():
         raise gr.Error("Enter input text first.")
 
+    t0 = time.perf_counter()
     cfg = copy.deepcopy(model_bundle["cfg"])
     cfg["inference"]["temperature"] = float(temperature)
     cfg["inference"]["top_k"] = int(top_k)
@@ -303,6 +344,30 @@ def generate_from_ui(
         "clean_output": bool(clean_output),
     }
     log_path = save_generation(model_bundle["experiment"], record)
+    latency_ms = (time.perf_counter() - t0) * 1000.0
+    toks = [t for t in output_text.split() if t]
+    uniq = len(set(toks)) / max(1, len(toks))
+    _mlflow_event(
+        run_name="space_inference",
+        params={
+            "experiment": model_bundle["experiment"],
+            "checkpoint": model_bundle["ckpt_path"],
+            "temperature": float(temperature),
+            "top_k": int(top_k),
+            "repetition_penalty": float(repetition_penalty),
+            "diversity_penalty": float(diversity_penalty),
+            "num_steps": int(num_steps),
+            "clean_output": bool(clean_output),
+        },
+        metrics={
+            "latency_ms": latency_ms,
+            "input_char_len": len(input_text.strip()),
+            "output_char_len": len(output_text),
+            "output_token_len": len(toks),
+            "output_unique_ratio": uniq,
+        },
+        tags={"source": "hf_space"},
+    )
     status = f"Inference done. Saved: `{log_path}`"
     return output_text, status, record
 
@@ -521,6 +586,7 @@ def _bg_worker(job_id: str, model_bundle, output_dir: str, input_text: str, task
     tasks = ["1", "2", "3", "4", "5"]
     failures = 0
     logs = []
+    run_start = time.perf_counter()
     _BG_JOBS[job_id].update({"state": "running", "progress": 0, "failures": 0, "updated": datetime.now().isoformat()})
     for idx, task in enumerate(tasks, start=1):
         _BG_JOBS[job_id].update(
@@ -545,6 +611,22 @@ def _bg_worker(job_id: str, model_bundle, output_dir: str, input_text: str, task
                 "updated": datetime.now().isoformat(),
             }
         )
+        _mlflow_event(
+            run_name=f"space_bg_task_{task}",
+            params={
+                "job_id": job_id,
+                "task": task,
+                "task4_phase": str(task4_phase),
+                "experiment": model_bundle.get("experiment", ""),
+            },
+            metrics={
+                "exit_code": float(code),
+                "used_bundled": 1.0 if used_bundled else 0.0,
+                "failures_so_far": float(failures),
+                "progress_pct": float(_BG_JOBS[job_id]["progress"]),
+            },
+            tags={"source": "hf_space", "mode": "background"},
+        )
     if failures:
         _bundle_task_outputs(model_bundle, output_dir)
     _BG_JOBS[job_id].update(
@@ -557,6 +639,20 @@ def _bg_worker(job_id: str, model_bundle, output_dir: str, input_text: str, task
             "updated": datetime.now().isoformat(),
         }
     )
+    _mlflow_event(
+        run_name="space_bg_run",
+        params={
+            "job_id": job_id,
+            "task4_phase": str(task4_phase),
+            "experiment": model_bundle.get("experiment", ""),
+            "output_dir": str(output_dir),
+        },
+        metrics={
+            "failures": float(failures),
+            "elapsed_s": (time.perf_counter() - run_start),
+        },
+        tags={"source": "hf_space", "mode": "background_summary"},
+    )
 
 
 def start_run_all_background(model_bundle, output_dir, input_text, task4_phase):
@@ -606,7 +702,9 @@ def run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task
 def run_single_task(model_bundle, task, output_dir, input_text, task4_phase):
     if not model_bundle:
         raise gr.Error("Load a model first.")
+    t0 = time.perf_counter()
     code, log, used_bundled = _run_analysis_cmd(task, model_bundle["ckpt_path"], output_dir, input_text, task4_phase)
+    elapsed = (time.perf_counter() - t0) * 1000.0
     if code != 0:
         _bundle_task_outputs(model_bundle, output_dir)
         log = f"{log}\n\n--- Live task analysis ---\n{_live_task_analysis(model_bundle, task, input_text)}"
@@ -618,6 +716,21 @@ def run_single_task(model_bundle, task, output_dir, input_text, task4_phase):
             status = f"Task {task} loaded from bundled analysis outputs + live analysis."
         else:
             status = f"Task {task} completed (exit={code})."
+    _mlflow_event(
+        run_name=f"space_task_{task}",
+        params={
+            "task": str(task),
+            "task4_phase": str(task4_phase),
+            "output_dir": str(output_dir),
+            "experiment": model_bundle.get("experiment", ""),
+        },
+        metrics={
+            "exit_code": float(code),
+            "elapsed_ms": elapsed,
+            "used_bundled": 1.0 if used_bundled else 0.0,
+        },
+        tags={"source": "hf_space", "mode": "single_task"},
+    )
     return status, log
 
 
@@ -692,6 +805,40 @@ def refresh_task_outputs(output_dir):
     )
 
 
+def _safe_refresh_task_outputs(output_dir):
+    try:
+        return refresh_task_outputs(output_dir)
+    except Exception as e:
+        err = f"Refresh error: {e}"
+        return (err, err, None, None, None, None, err, None, err, None)
+
+
+def _safe_start_run_all_background(model_bundle, output_dir, input_text, task4_phase, current_job_id):
+    try:
+        status, log, job_id = start_run_all_background(model_bundle, output_dir, input_text, task4_phase)
+        return status, log, job_id
+    except Exception as e:
+        return f"Background start failed: {e}", f"Background start failed: {e}", current_job_id
+
+
+def _safe_poll_run_all_background(job_id, output_dir):
+    try:
+        return poll_run_all_background(job_id, output_dir)
+    except Exception as e:
+        err = f"Track error: {e}"
+        out = _safe_refresh_task_outputs(output_dir)
+        return err, err, *out
+
+
+def _safe_run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task4_phase):
+    try:
+        return run_single_task_and_refresh(model_bundle, task, output_dir, input_text, task4_phase)
+    except Exception as e:
+        err = f"Task {task} failed: {e}"
+        out = _safe_refresh_task_outputs(output_dir)
+        return err, err, *out
+
+
 CUSTOM_CSS = """
 :root {
   --bg1: #f5fbff;
@@ -886,7 +1033,7 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
     )
 
     run_single_btn.click(
-        fn=run_single_task_and_refresh,
+        fn=_safe_run_single_task_and_refresh,
         inputs=[model_state, task_choice, analysis_output_dir, analysis_input, task4_phase],
         outputs=[
             task_run_status,
@@ -904,12 +1051,12 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
         ],
     )
     run_all_btn.click(
-        fn=start_run_all_background,
-        inputs=[model_state, analysis_output_dir, analysis_input, task4_phase],
+        fn=_safe_start_run_all_background,
+        inputs=[model_state, analysis_output_dir, analysis_input, task4_phase, bg_job_state],
         outputs=[task_run_status, task_run_log, bg_job_state],
     )
     track_bg_btn.click(
-        fn=poll_run_all_background,
+        fn=_safe_poll_run_all_background,
         inputs=[bg_job_state, analysis_output_dir],
         outputs=[
             task_run_status,
@@ -927,7 +1074,7 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
         ],
     )
     refresh_outputs_btn.click(
-        fn=refresh_task_outputs,
+        fn=_safe_refresh_task_outputs,
         inputs=[analysis_output_dir],
         outputs=[
             task1_box,
@@ -947,7 +1094,7 @@ with gr.Blocks(title="Sanskrit Diffusion Client Demo", css=CUSTOM_CSS) as demo:
         outputs=[model_state, load_status, model_info, num_steps, analysis_output_dir],
     )
     demo.load(
-        fn=refresh_task_outputs,
+        fn=_safe_refresh_task_outputs,
         inputs=[analysis_output_dir],
         outputs=[
             task1_box,

requirements.txt

@@ -4,3 +4,4 @@ numpy>=1.24
 tqdm>=4.66
 huggingface_hub>=0.30,<1.0
 tokenizers>=0.15
+mlflow-skinny>=2.16.0