155 lines
4.8 KiB
Python
155 lines
4.8 KiB
Python
import os
|
|
import ast
|
|
from typing import Optional, Tuple, Union
|
|
import random
|
|
import numpy as np
|
|
import torch
|
|
|
|
|
|
DTYPE_MAPPING = {
|
|
"float32": torch.float32,
|
|
"float16": torch.float16,
|
|
"bfloat16": torch.bfloat16,
|
|
"float64": torch.float64,
|
|
"uint8": torch.uint8,
|
|
"int8": torch.int8,
|
|
"int16": torch.int16,
|
|
"int32": torch.int32,
|
|
"int64": torch.int64,
|
|
"bool": torch.bool
|
|
}
|
|
|
|
def str_to_number(s: str) -> Union[int, float]:
|
|
"""
|
|
Converts a string representation of a number to an integer or float.
|
|
|
|
Args:
|
|
s (str): The string to convert.
|
|
|
|
Returns:
|
|
Union[int, float]: The converted number, either as an `int` if no decimal point is present or as a `float` otherwise.
|
|
|
|
Raises:
|
|
ValueError: If the input string is empty.
|
|
"""
|
|
if s:
|
|
if "." in s:
|
|
return float(s)
|
|
return int(s)
|
|
raise ValueError("Input string cannot be empty.")
|
|
|
|
|
|
def str_to_number_with_default(s: str, n: Union[int, float]) -> Union[int, float]:
|
|
"""
|
|
Converts a string representation of a number to an integer or float, with a fallback default value.
|
|
|
|
Args:
|
|
s (str): The string to convert.
|
|
n (Union[int, float]): The default value to return if `s` is empty.
|
|
|
|
Returns:
|
|
Union[int, float]: The converted number if `s` is non-empty, otherwise `n`.
|
|
"""
|
|
return str_to_number(s) if s else n
|
|
|
|
def get_resource_full_path(resource_name: str, resource_path: str) -> str:
|
|
"""
|
|
Resolves the absolute path to a resource file within the "resources" directory and
|
|
ensures that the directory exists.
|
|
|
|
Args:
|
|
resource_path (str): The path to the resource file. If the path is a relative path,
|
|
the path is considered relative to the `resource` directory of this extension.
|
|
|
|
Returns:
|
|
str: The absolute path to the resource file.
|
|
"""
|
|
if os.path.isabs(resource_path):
|
|
resource_full_path = resource_path
|
|
else:
|
|
resource_full_path = os.path.realpath(
|
|
os.path.join(os.path.dirname(__file__), "..", "..", resource_name, resource_path))
|
|
resource_dir = os.path.dirname(resource_full_path)
|
|
if os.path.exists(resource_dir) is False:
|
|
os.makedirs(resource_dir)
|
|
return resource_full_path
|
|
|
|
|
|
def get_model_full_path(model_path: str) -> str:
|
|
"""
|
|
Resolves the absolute path to a model file within the "models" directory and
|
|
ensures that the directory exists.
|
|
|
|
Args:
|
|
model_path (str): The path to the model file. If the path is a relative path,
|
|
the path is considered relative to the `models` directory of this extension.
|
|
|
|
Returns:
|
|
str: The absolute path to the model file.
|
|
"""
|
|
return get_resource_full_path("models", model_path)
|
|
|
|
|
|
def get_dataset_full_path(dataset_path: str) -> str:
|
|
"""
|
|
Resolves the absolute path to the datasete within the "datasets" directory and
|
|
ensures that the directory exists.
|
|
|
|
Args:
|
|
dataset_path (str): The path to the dataset file. If the path is a relative path,
|
|
the path is considered relative to the `datasets` directory of this extension.
|
|
|
|
Returns:
|
|
str: The absolute path to the dataset file.
|
|
"""
|
|
return get_resource_full_path("datasets", dataset_path)
|
|
|
|
|
|
def str_to_dim(dim: str) -> Optional[Union[int, Tuple[int, ...]]]:
|
|
"""
|
|
Converts a string representation of a dimension into an integer or a tuple of integers.
|
|
|
|
Args:
|
|
dim (str): A string representing an integer or a tuple of integers.
|
|
|
|
Returns:
|
|
Optional[Union[int, Tuple[int, ...]]]:
|
|
- An integer if the input represents a single integer.
|
|
- A tuple of integers if the input represents a list or tuple.
|
|
- None if the input is empty.
|
|
|
|
Raises:
|
|
ValueError: If the parsed value is not an integer or a tuple of integers.
|
|
"""
|
|
if dim:
|
|
# Convert string to an integer, list, or tuple of integers
|
|
parsed_dim = ast.literal_eval(dim)
|
|
|
|
# Convert list to tuple
|
|
if isinstance(parsed_dim, list):
|
|
parsed_dim = tuple(parsed_dim)
|
|
|
|
# Ensure the parsed dim is valid
|
|
if not isinstance(parsed_dim, (int, tuple)) or (
|
|
isinstance(parsed_dim, tuple) and not all(isinstance(i, int) for i in parsed_dim)
|
|
):
|
|
raise ValueError(f"Invalid format for dim: {parsed_dim}. Must be an integer or a tuple of integers.")
|
|
|
|
return parsed_dim
|
|
|
|
return None
|
|
|
|
def set_seed(seed=42) -> None:
|
|
"""
|
|
Sets seed for random number generation.
|
|
|
|
Args:
|
|
seed (int): Initial seed.
|
|
"""
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
random.seed(seed)
|
|
if torch.cuda.is_available():
|
|
torch.cuda.manual_seed_all(seed)
|
|
torch.backends.cudnn.deterministic = True
|
|
torch.backends.cudnn.benchmark = False |