Upload diffusion_llm/training/__init__.py with huggingface_hub
Browse files
diffusion_llm/training/__init__.py
CHANGED
|
@@ -3,5 +3,6 @@
|
|
| 3 |
from diffusion_llm.training.trainer import AamTrainer
|
| 4 |
from diffusion_llm.training.dataset import GraphNarrativeDataset
|
| 5 |
from diffusion_llm.training.losses import DiffusionLoss, compute_loss
|
|
|
|
| 6 |
|
| 7 |
-
__all__ = ["AamTrainer", "GraphNarrativeDataset", "DiffusionLoss", "compute_loss"]
|
|
|
|
| 3 |
from diffusion_llm.training.trainer import AamTrainer
|
| 4 |
from diffusion_llm.training.dataset import GraphNarrativeDataset
|
| 5 |
from diffusion_llm.training.losses import DiffusionLoss, compute_loss
|
| 6 |
+
from diffusion_llm.training.llm_jepa import JEPAPredictor, JEPAConfig, JEPATrainer
|
| 7 |
|
| 8 |
+
__all__ = ["AamTrainer", "GraphNarrativeDataset", "DiffusionLoss", "compute_loss", "JEPAPredictor", "JEPAConfig", "JEPATrainer"]
|