Commit
·
e414316
1
Parent(s):
03d21c5
minor formatting
Browse files- modeling_qwen2nsa.py +2 -0
modeling_qwen2nsa.py
CHANGED
|
@@ -25,6 +25,7 @@ except ImportError:
|
|
| 25 |
|
| 26 |
from native_sparse_attention_pytorch import SparseAttention
|
| 27 |
|
|
|
|
| 28 |
class Qwen2RMSNorm(nn.Module):
|
| 29 |
def __init__(self, hidden_size: int, eps: float = 1e-6):
|
| 30 |
super().__init__()
|
|
@@ -233,6 +234,7 @@ class Qwen2Model(nn.Module):
|
|
| 233 |
out["past_key_values"] = next_cache
|
| 234 |
return out
|
| 235 |
|
|
|
|
| 236 |
class Qwen2NsaForCausalLM(PreTrainedModel, GenerationMixin):
|
| 237 |
config_class = Qwen2NsaConfig
|
| 238 |
base_model_prefix = "model"
|
|
|
|
| 25 |
|
| 26 |
from native_sparse_attention_pytorch import SparseAttention
|
| 27 |
|
| 28 |
+
|
| 29 |
class Qwen2RMSNorm(nn.Module):
|
| 30 |
def __init__(self, hidden_size: int, eps: float = 1e-6):
|
| 31 |
super().__init__()
|
|
|
|
| 234 |
out["past_key_values"] = next_cache
|
| 235 |
return out
|
| 236 |
|
| 237 |
+
|
| 238 |
class Qwen2NsaForCausalLM(PreTrainedModel, GenerationMixin):
|
| 239 |
config_class = Qwen2NsaConfig
|
| 240 |
base_model_prefix = "model"
|