♻️ match updated module name

Signed-off-by: peter szemraj <peterszemraj@gmail.com>

Author: pszemraj

samba_pytorch/config.py

@@ -7,11 +7,12 @@
 from dataclasses import dataclass
 from typing import Any, Literal, Optional, Type
 
-import lit_gpt.model
 import torch
-from lit_gpt.utils import find_multiple
 from typing_extensions import Self
 
+import samba_pytorch.samba
+from samba_pytorch.utils import find_multiple
+
 
 @dataclass
 class Config:
@@ -101,17 +102,17 @@ def from_name(cls, name: str, **kwargs: Any) -> Self:
     @property
     def mlp_class(self) -> Type:
         # `self._mlp_class` cannot be the type to keep the config json serializable
-        return getattr(lit_gpt.model, self._mlp_class)
+        return getattr(samba_pytorch.samba, self._mlp_class)
 
     @property
     def norm_class(self) -> Type:
         # `self._norm_class` cannot be the type to keep the config json serializable
         if self._norm_class == "RMSNorm":
-            from lit_gpt.rmsnorm import RMSNorm
+            from samba_pytorch.modules.rmsnorm import RMSNorm
 
             return RMSNorm
         elif self._norm_class == "FusedRMSNorm":
-            from lit_gpt.rmsnorm import FusedRMSNorm
+            from samba_pytorch.modules.rmsnorm import FusedRMSNorm
 
             return FusedRMSNorm
         return getattr(torch.nn, self._norm_class)

samba_pytorch/samba.py

@@ -10,23 +10,22 @@
 
 import torch
 import torch.nn as nn
-from lit_gpt.config import Config
 from torch import Tensor
 from typing_extensions import Self
 from xformers.ops import SwiGLU
 
-from .fused_rotary_embedding import apply_rotary_emb_func
-from .mamba_simple import Mamba
-
 try:
     from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn
 except ImportError:
     RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None
 from causal_conv1d import causal_conv1d_fn
 from einops import rearrange
 
-from .gla import GatedLinearAttention
-from .multiscale_retention import MultiScaleRetention
+from samba_pytorch.config import Config
+from samba_pytorch.modules.fused_rotary_embedding import apply_rotary_emb_func
+from samba_pytorch.modules.gla import GatedLinearAttention
+from samba_pytorch.modules.mamba_simple import Mamba
+from samba_pytorch.modules.multiscale_retention import MultiScaleRetention
 
 RoPECache = Tuple[torch.Tensor, torch.Tensor]
 KVCache = Tuple[torch.Tensor, torch.Tensor]