Spaces:
Sleeping
Sleeping
File size: 1,747 Bytes
c1b1f3a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 | import os
import sys
import platform
import mlflow
from typing import Dict, Any, Optional
def setup_mlflow(config: dict) -> str:
"""Sets up local MLflow tracking server URI and artifact directory."""
tracking_uri = config.get("tracking", {}).get("tracking_uri", "sqlite:///mlflow.db")
mlflow.set_tracking_uri(tracking_uri)
# Disable annoying hints
os.environ["MLFLOW_DISABLE_AGENT_HINT"] = "1"
# Pre-create standard experiments
experiments = config.get("tracking", {}).get("experiments", {
"ocr_comparison": "ocr_comparison",
"doc_classification": "doc_classification",
"field_extraction": "field_extraction",
"robustness_tests": "robustness_tests",
"end_to_end": "end_to_end",
"threshold_analysis": "threshold_analysis"
})
for exp_key, exp_name in experiments.items():
try:
mlflow.create_experiment(exp_name)
except Exception:
pass # Already exists
return tracking_uri
def get_standard_run_tags(stage: str, model_family: str, purpose: str = "baseline") -> Dict[str, str]:
"""Generates standard run tags required by project specs."""
return {
"stage": stage,
"model_family": model_family,
"dataset_version": "1.0.0",
"author": "doc_intel_agent",
"purpose": purpose,
"os": platform.system(),
"python_version": platform.python_version()
}
def get_standard_run_params(seed: int = 42, **kwargs) -> Dict[str, Any]:
"""Generates standard parameter dictionary."""
params = {
"seed": seed,
"hardware": "cuda" if False else "cpu",
"python": platform.python_version()
}
params.update(kwargs)
return params
|