Coverage for dantinox / __init__.py: 87%

1""" 

2DantinoX — a decoder-only Transformer library built in JAX and Flax NNX. 

3 

4Quick start 

5----------- 

6>>> from dantinox import Config, Transformer, Trainer, Generator 

7>>> 

8>>> config = Config.from_yaml("configs/default_config.yaml") 

9>>> trainer = Trainer(config) 

10>>> run_dir = trainer.fit("data/corpus.txt") 

11>>> 

12>>> gen = Generator(run_dir) 

13>>> print(gen.generate("Nel mezzo del cammin ")) 

14""" 

15 

16from importlib.metadata import PackageNotFoundError, version 

17 

18try: 

19 __version__: str = version("dantinox") 

20except PackageNotFoundError: 

21 __version__ = "0.0.0.dev" 

22 

23from core.config import Config 

24from core.generation import decode, generate 

25from core.model import Transformer 

26from dantinox.bench import BenchmarkRunner 

27from dantinox.exceptions import ( 

28 BenchmarkError, 

29 CheckpointError, 

30 ConfigError, 

31 DantinoXError, 

32 PlotError, 

33) 

34from dantinox.generator import Generator 

35from dantinox.hub import pull, push, resolve_checkpoint 

36from dantinox.plotting import Plotter 

37from dantinox.trainer import Trainer 

38 

39__all__ = [ 

40 # version 

41 "__version__", 

42 # core 

43 "Config", 

44 "Transformer", 

45 "generate", 

46 "decode", 

47 # high-level API 

48 "Trainer", 

49 "Generator", 

50 "BenchmarkRunner", 

51 "Plotter", 

52 "push", 

53 "pull", 

54 "resolve_checkpoint", 

55 # exceptions 

56 "DantinoXError", 

57 "ConfigError", 

58 "CheckpointError", 

59 "BenchmarkError", 

60 "PlotError", 

61]