Spaces:
Running
Running
| #!/usr/bin/env python3 | |
| # utils.py - Common utilities for XSimGCL project | |
| """ | |
| Shared helper functions for data conversion, hardware setup, | |
| configuration management, checkpoint handling. | |
| FIXED: resolve_config_path now checks script directory FIRST before going up | |
| ADDED: load_active_config() — czyta train_dataset.yaml i zwraca (config_path, dataset_type) | |
| """ | |
| import torch | |
| import numpy as np | |
| import scipy.sparse as sp | |
| import yaml | |
| import os | |
| import sys | |
| from typing import Dict, Optional, List, Tuple | |
| # ============================================================================ | |
| # SECTION 1: Data Conversion | |
| # ============================================================================ | |
| def scipy_to_torch_sparse(mat): | |
| """ | |
| Convert scipy sparse matrix to torch sparse tensor. | |
| Używa formatu CSR (Compressed Sparse Row) jeśli PyTorch >= 1.11. | |
| CSR jest cache-friendly dla torch.sparse.mm na CUDA — szybszy forward pass GCN. | |
| Fallback do COO dla starszych wersji PyTorch. | |
| Args: | |
| mat: Scipy sparse matrix or dense array | |
| Returns: | |
| torch.sparse.Tensor (CSR preferred) or original input if not sparse | |
| """ | |
| if not sp.issparse(mat): | |
| return mat | |
| csr = mat.tocsr() | |
| crow_indices = torch.from_numpy(csr.indptr.astype(np.int64)) | |
| col_indices = torch.from_numpy(csr.indices.astype(np.int64)) | |
| values = torch.from_numpy(csr.data.astype(np.float32)) | |
| try: | |
| return torch.sparse_csr_tensor( | |
| crow_indices, col_indices, values, torch.Size(csr.shape) | |
| ) | |
| except Exception: | |
| # Fallback: COO dla PyTorch < 1.11 | |
| coo = mat.tocoo() | |
| indices = torch.from_numpy(np.vstack((coo.row, coo.col)).astype(np.int64)) | |
| values_coo = torch.from_numpy(coo.data.astype(np.float32)) | |
| return torch.sparse_coo_tensor(indices, values_coo, torch.Size(coo.shape)).coalesce() | |
| # ============================================================================ | |
| # SECTION 2: Hardware Setup | |
| # ============================================================================ | |
| def setup_device(device_arg: str) -> str: | |
| """ | |
| Auto-detect or select compute device. | |
| Args: | |
| device_arg: 'auto', 'cuda', or 'cpu' | |
| Returns: | |
| Selected device string ('cuda' or 'cpu') | |
| """ | |
| if device_arg == 'auto': | |
| if torch.cuda.is_available(): | |
| print(f"[Device] CUDA: {torch.cuda.get_device_name(0)}") | |
| return 'cuda' | |
| print("[Device] Using CPU") | |
| return 'cpu' | |
| if device_arg == 'cuda' and not torch.cuda.is_available(): | |
| print("[Device] CUDA requested but not available → falling back to CPU") | |
| return 'cpu' | |
| return device_arg | |
| # ============================================================================ | |
| # SECTION 3: Config Management | |
| # ============================================================================ | |
| def load_yaml_config(path: str, required: bool = True) -> Optional[Dict]: | |
| """ | |
| Load YAML configuration file. | |
| Args: | |
| path: Path to YAML file | |
| required: If True, exit on missing file | |
| Returns: | |
| Configuration dictionary or None | |
| """ | |
| if not os.path.exists(path): | |
| if required: | |
| print(f"[ERROR] Missing configuration file: {path}") | |
| sys.exit(1) | |
| return None | |
| try: | |
| with open(path, 'r') as f: | |
| return yaml.safe_load(f) | |
| except Exception as e: | |
| print(f"[ERROR] Failed to load YAML: {e}") | |
| if required: | |
| sys.exit(1) | |
| return None | |
| def load_active_config(script_dir: Optional[str] = None) -> Tuple[str, str]: | |
| """ | |
| Czyta train_dataset.yaml i zwraca (config_path, dataset_type). | |
| Pozwala wszystkim skryptom (train_model, png, test_recommend, grid_train) | |
| automatycznie wiedzieć którego datasetu i której konfiguracji używać | |
| bez ręcznego podawania argumentów. | |
| train_dataset.yaml: | |
| dataset: "f95" # lub "steam" | |
| Returns: | |
| (config_path, dataset_type) gdzie dataset_type in {'f95', 'steam'} | |
| """ | |
| if script_dir is None: | |
| script_dir = get_script_directory() | |
| td_path = os.path.join(script_dir, 'train_dataset.yaml') | |
| if not os.path.exists(td_path): | |
| # Fallback: szukaj w katalogu roboczym | |
| td_path = os.path.join(os.getcwd(), 'train_dataset.yaml') | |
| td = load_yaml_config(td_path, required=False) | |
| if not td: | |
| # Brak pliku → domyślnie f95 | |
| print(f" [load_active_config] train_dataset.yaml not found, defaulting to f95") | |
| config_path = os.path.join(script_dir, 'config.yaml') | |
| return config_path, 'f95' | |
| dataset_type = str(td.get('dataset', 'f95')).lower().strip() | |
| if dataset_type == 'steam': | |
| config_path = os.path.join(script_dir, 'config_steam.yaml') | |
| else: | |
| dataset_type = 'f95' | |
| config_path = os.path.join(script_dir, 'config.yaml') | |
| print(f" [Dataset] train_dataset.yaml → dataset={dataset_type}, config={os.path.basename(config_path)}") | |
| return config_path, dataset_type | |
| def load_dataset_class(dataset_type: str): | |
| """ | |
| Zwraca klasę datasetu na podstawie dataset_type. | |
| Args: | |
| dataset_type: 'f95' lub 'steam' | |
| Returns: | |
| Klasa datasetu (F95Dataset lub SteamDataset) | |
| """ | |
| if dataset_type == 'steam': | |
| from game_dataset_steam import SteamDataset | |
| return SteamDataset | |
| else: | |
| from game_dataset import F95Dataset | |
| return F95Dataset | |
| # ============================================================================ | |
| # SECTION 4: Checkpoint Management | |
| # ============================================================================ | |
| def find_checkpoint(base_dir: str, checkpoint_patterns: List[str]) -> Optional[str]: | |
| """ | |
| Search for checkpoint directory matching patterns. | |
| Args: | |
| base_dir: Base directory to search in | |
| checkpoint_patterns: List of subdirectory names to try | |
| Returns: | |
| Path to checkpoint directory or None | |
| """ | |
| for pattern in checkpoint_patterns: | |
| checkpoint_dir = os.path.join(base_dir, pattern) | |
| if os.path.exists(os.path.join(checkpoint_dir, 'best_model.pt')): | |
| return checkpoint_dir | |
| return None | |
| # ============================================================================ | |
| # SECTION 5: Path Utilities (FIXED) | |
| # ============================================================================ | |
| def get_script_directory() -> str: | |
| """ | |
| Get directory where the calling script is located. | |
| Returns: | |
| Absolute path to script directory | |
| """ | |
| if hasattr(sys.modules['__main__'], '__file__'): | |
| script_path = os.path.abspath(sys.modules['__main__'].__file__) | |
| return os.path.dirname(script_path) | |
| return os.path.abspath(os.getcwd()) | |
| def get_project_root() -> str: | |
| """ | |
| Get project root directory (where config.yaml is expected). | |
| FIXED: Now searches from SCRIPT directory first, then goes up. | |
| Returns: | |
| Absolute path to project root | |
| """ | |
| current = get_script_directory() | |
| if os.path.exists(os.path.join(current, 'config.yaml')): | |
| return current | |
| search_path = current | |
| while search_path != '/': | |
| if os.path.exists(os.path.join(search_path, 'config.yaml')): | |
| return search_path | |
| search_path = os.path.dirname(search_path) | |
| return current | |
| def resolve_config_path(custom_path: Optional[str] = None) -> str: | |
| """ | |
| Resolve configuration file path. | |
| UWAGA: Ta funkcja zwraca config.yaml (F95) bez uwzględnienia train_dataset.yaml. | |
| Używaj load_active_config() jeśli chcesz automatyczną selekcję datasetu. | |
| Args: | |
| custom_path: Custom path to config.yaml (optional) | |
| Returns: | |
| Resolved absolute path to config.yaml | |
| """ | |
| if custom_path: | |
| return os.path.abspath(custom_path) | |
| project_root = get_project_root() | |
| return os.path.join(project_root, 'config.yaml') | |