Spaces:
Sleeping
Sleeping
| 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 | |