File size: 2,788 Bytes
3d79eb3 |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 |
# Copyright (c) Meta Platforms, Inc. and affiliates
# All rights reserved.
#
#
import ctypes
from abc import abstractmethod
from pathlib import Path
from typing import (
Any,
Dict,
Iterable,
Optional,
Protocol,
Sized,
Type,
TypeVar,
Union,
runtime_checkable,
)
import torch
from omegaconf import DictConfig, OmegaConf
root_working_dir = Path(__file__).parent.parent.parent
def set_mkl_num_threads():
"""Setting mkl num threads to 1, so that we don't get thread explosion."""
mkl_rt = ctypes.CDLL("libmkl_rt.so")
mkl_rt.mkl_set_num_threads(ctypes.byref(ctypes.c_int(1)))
def working_dir_resolver(p: str):
"""The omegaconf resolver that translates a relative path to the absolute path"""
return "file://" + str(root_working_dir.joinpath(p).resolve())
def setup_conf():
"""Register the common Hydra config groups used in LCM (for now only the launcher)"""
from stopes.pipelines import config_registry # noqa
recipe_root = Path(__file__).parent.parent.parent / "recipes"
config_registry["lcm-common"] = "file://" + str((recipe_root / "common").resolve())
config_registry["lcm-root"] = "file://" + str(recipe_root.resolve())
# Register omegaconf resovlers
OmegaConf.register_new_resolver("realpath", working_dir_resolver, replace=True)
def torch_type(
dtype: Optional[Union[str, torch.dtype]] = None,
) -> Optional[torch.dtype]:
# Convert dtyp string from the checkpoint to torch.dtype
# https://github.com/pytorch/pytorch/issues/40471
if dtype is None:
return None
if isinstance(dtype, torch.dtype):
return dtype
_dtype = eval(dtype) # type: ignore
assert isinstance(_dtype, torch.dtype), f"Invalid dtype value: {dtype}"
return _dtype
@runtime_checkable
class Batched(Sized, Protocol):
"""Abstract class for batched data"""
@abstractmethod
def __getitem__(self, i: int) -> Any: ...
T = TypeVar("T")
def promote_config(config: Union[T, DictConfig, Dict], config_cls: Type[T]) -> T:
if isinstance(config, (Dict, DictConfig)):
import dacite
if isinstance(config, DictConfig):
config = OmegaConf.to_container(config) # type: ignore
return dacite.from_dict(
data_class=config_cls,
data=config, # type: ignore
config=dacite.Config(cast=[Path]), # type: ignore
)
else:
assert isinstance(config, config_cls), f"Unknown config type: {type(config)}"
return config
def batched(inputs: Iterable, batch_size=10000) -> Iterable:
batch = []
for line in inputs:
batch.append(line)
if len(batch) == batch_size:
yield batch
batch = []
if len(batch) > 0:
yield batch
|