Create utils.py

utils.py

@@ -0,0 +1,193 @@
+"""
+OmniCoreX Utilities Module
+
+Helper functions, logging setup, configuration parsing,
+and common utilities used throughout the OmniCoreX system.
+
+Features:
+- Robust logging setup with configurable formats and levels.
+- Configuration loader supporting YAML and JSON with overrides.
+- Seed setting for reproducibility.
+- Timing and benchmarking decorators.
+- Various small utilities for system use.
+"""
+
+import os
+import sys
+import yaml
+import json
+import logging
+import random
+import time
+import numpy as np
+import torch
+
+# ----------------------- Logging Setup ----------------------- #
+
+def setup_logging(log_level=logging.INFO, log_file: str = None) -> logging.Logger:
+    """
+    Sets up a logger with console and optional file handlers.
+
+    Args:
+        log_level: Logging level (e.g., logging.INFO).
+        log_file: Optional path to log file.
+
+    Returns:
+        Configured logger instance.
+    """
+    logger = logging.getLogger("OmniCoreX")
+    logger.setLevel(log_level)
+    formatter = logging.Formatter("%(asctime)s | %(levelname)s | %(name)s | %(message)s")
+
+    # Remove existing handlers
+    for handler in logger.handlers[:]:
+        logger.removeHandler(handler)
+
+    # Console handler
+    ch = logging.StreamHandler(sys.stdout)
+    ch.setLevel(log_level)
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)
+
+    # File handler if specified
+    if log_file:
+        fh = logging.FileHandler(log_file)
+        fh.setLevel(log_level)
+        fh.setFormatter(formatter)
+        logger.addHandler(fh)
+
+    return logger
+
+# Global logger instance
+logger = setup_logging()
+
+# ----------------------- Configuration Loading ----------------------- #
+
+def load_config_file(config_path: str) -> dict:
+    """
+    Loads a YAML or JSON configuration file.
+
+    Args:
+        config_path: Path to the config file.
+
+    Returns:
+        Dictionary of configuration parameters.
+    """
+    if not os.path.isfile(config_path):
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+
+    ext = os.path.splitext(config_path)[1].lower()
+    with open(config_path, "r", encoding="utf-8") as f:
+        if ext in [".yaml", ".yml"]:
+            cfg = yaml.safe_load(f)
+        elif ext == ".json":
+            cfg = json.load(f)
+        else:
+            raise ValueError(f"Unsupported config format: {ext}")
+
+    return cfg
+
+def merge_dicts(base: dict, override: dict) -> dict:
+    """
+    Deep merges two dictionaries, with the override taking precedence.
+
+    Args:
+        base: Base dictionary.
+        override: Dictionary with override values.
+
+    Returns:
+        Merged dictionary.
+    """
+    result = base.copy()
+    for k, v in override.items():
+        if k in result and isinstance(result[k], dict) and isinstance(v, dict):
+            result[k] = merge_dicts(result[k], v)
+        else:
+            result[k] = v
+    return result
+
+# ----------------------- Seed Setting ----------------------- #
+
+def set_seed(seed: int = 42):
+    """
+    Set seed for reproducibility across random, numpy and torch.
+
+    Args:
+        seed: Integer seed value.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    logger.info(f"Random seed set to {seed}")
+
+# ----------------------- Timing Utilities ----------------------- #
+
+def timeit(func):
+    """
+    Decorator to measure and log function execution time.
+
+    Usage:
+        @timeit
+        def my_function(...):
+            ...
+    """
+    def wrapper(*args, **kwargs):
+        start = time.time()
+        result = func(*args, **kwargs)
+        end = time.time()
+        logger.info(f"Function {func.__name__!r} executed in {(end - start):.4f}s")
+        return result
+    return wrapper
+
+# ----------------------- Other Utility Functions ----------------------- #
+
+def ensure_dir(dirname: str):
+    """
+    Creates directory if it does not exist.
+
+    Args:
+        dirname: Directory path to create.
+    """
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)
+        logger.debug(f"Directory created: {dirname}")
+
+def to_device(batch: dict, device: torch.device) -> dict:
+    """
+    Moves all tensor elements in batch dict to specified device.
+
+    Args:
+        batch: Dictionary with tensors.
+        device: Target torch device.
+
+    Returns:
+        Dictionary with tensors on device.
+    """
+    return {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+
+if __name__ == "__main__":
+    # Demo usage of utilities
+
+    set_seed(1234)
+    logger.info("This is a test log message.")
+
+    # Create dummy config files and test merging
+    base_cfg = {"model": {"layers": 12, "embed_dim": 256}, "training": {"batch_size": 32}}
+    override_cfg = {"model": {"layers": 24}, "training": {"learning_rate": 0.001}}
+
+    merged_cfg = merge_dicts(base_cfg, override_cfg)
+    logger.info(f"Merged config: {merged_cfg}")
+
+    # Test directory creation
+    test_dir = "./tmp_test_dir"
+    ensure_dir(test_dir)
+
+    # Test timing decorator
+    @timeit
+    def dummy_work():
+        import time; time.sleep(0.5)
+
+    dummy_work()
+