Spaces:
Runtime error
Runtime error
from dataclasses import dataclass | |
from typing import TypeVar, Generic, Type, Optional | |
from functools import wraps | |
import time | |
import random | |
import torch as T | |
import torch.nn as nn | |
# @TODO: remove si_module from codebase | |
# we use this in our research codebase to make modules from callable configs | |
si_module_TpV = TypeVar('si_module_TpV') | |
def si_module(cls: Type[si_module_TpV]) -> Type[si_module_TpV]: | |
if not hasattr(cls, 'Config') or not isinstance(cls.Config, type): | |
class Config: | |
pass | |
cls.Config = Config | |
cls.Config = dataclass(cls.Config) | |
class ConfigWrapper(cls.Config, Generic[si_module_TpV]): | |
def __call__(self, *args, **kwargs) -> si_module_TpV: | |
if len(kwargs) > 0: | |
config_dict = {field.name: getattr(self, field.name) for field in self.__dataclass_fields__.values()} | |
config_dict.update(kwargs) | |
new_config = type(self)(**config_dict) | |
return cls(new_config) | |
else: | |
return cls(self, *args) | |
ConfigWrapper.__module__ = cls.__module__ | |
ConfigWrapper.__name__ = f"{cls.__name__}Config" | |
ConfigWrapper.__qualname__ = f"{cls.__qualname__}.Config" | |
cls.Config = ConfigWrapper | |
original_init = cls.__init__ | |
def new_init(self, *args, **kwargs): | |
self.c = next((arg for arg in args if isinstance(arg, cls.Config)), None) or next((arg for arg in kwargs.values() if isinstance(arg, cls.Config)), None) | |
original_init(self, *args, **kwargs) | |
self.register_buffer('_device_tracker', T.Tensor(), persistent=False) | |
cls.__init__ = new_init | |
def device(self): | |
return self._device_tracker.device | |
def dtype(self): | |
return self._device_tracker.dtype | |
cls.device = device | |
cls.dtype = dtype | |
return cls | |
def get_activation(nonlinear_activation, nonlinear_activation_params={}): | |
if hasattr(nn, nonlinear_activation): | |
return getattr(nn, nonlinear_activation)(**nonlinear_activation_params) | |
else: | |
raise NotImplementedError(f"Activation {nonlinear_activation} not found in torch.nn") | |
def exists(v): | |
return v is not None | |
def isnt(v): | |
return not exists(v) | |
def truthyexists(v): | |
return exists(v) and v is not False | |
def truthyattr(obj, attr): | |
return hasattr(obj, attr) and truthyexists(getattr(obj, attr)) | |
defaultT = TypeVar('defaultT') | |
def default(*args: Optional[defaultT]) -> Optional[defaultT]: | |
for arg in args: | |
if exists(arg): | |
return arg | |
return None | |
def maybe(fn): | |
def inner(x, *args, **kwargs): | |
if not exists(x): | |
return x | |
return fn(x, *args, **kwargs) | |
return inner | |