Upload folder using huggingface_hub

config.json

@@ -8,7 +8,7 @@
     "AutoModelForCausalLM": "lance_ai_model.LanceAI"
   },
   "bos_token_id": 151643,
-  "dtype": "bfloat16",
+  "dtype": "float32",
   "eos_token_id": 151643,
   "head_dim": 64,
   "hidden_act": "silu",

lance_ai_model.py

@@ -60,6 +60,7 @@ class LanceAIConfig(PretrainedConfig):
         self.bos_token_id = bos_token_id
         self.eos_token_id = eos_token_id
 
+
 class LanceAIRMSNorm(nn.Module):
     def __init__(self, hidden_size, eps=1e-6):
         super().__init__()
@@ -73,6 +74,7 @@ class LanceAIRMSNorm(nn.Module):
         hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
         return self.weight * hidden_states.to(input_dtype)
 
+
 class LanceAIRotaryEmbedding(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -92,6 +94,7 @@ class LanceAIRotaryEmbedding(nn.Module):
         sin = emb.sin()
         return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)
 
+
 def rotate_half(x):
     x1 = x[..., :x.shape[-1] // 2]
     x2 = x[..., x.shape[-1] // 2:]
@@ -113,6 +116,7 @@ def repeat_kv(hidden_states, n_rep):
     hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)
     return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
 
+
 def eager_attention_forward(
     module: nn.Module,
     query: torch.Tensor,
@@ -138,6 +142,7 @@ def eager_attention_forward(
 
     return attn_output, attn_weights
 
+
 class LanceAIAttention(nn.Module):
     def __init__(self, config, layer_idx):
         super().__init__()
@@ -200,6 +205,7 @@ class LanceAIAttention(nn.Module):
         attn_output = self.o_proj(attn_output)
         return attn_output, attn_weights
 
+
 class LanceAIMLP(nn.Module):
     def __init__(self, config):
         super().__init__()
@@ -212,6 +218,7 @@ class LanceAIMLP(nn.Module):
     def forward(self, x):
         return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x))
 
+
 class LanceAIDecoderLayer(GradientCheckpointingLayer):
     def __init__(self, config, layer_idx):
         super().__init__()
@@ -249,6 +256,7 @@ class LanceAIDecoderLayer(GradientCheckpointingLayer):
 
         return hidden_states
 
+
 class LanceAIPreTrainedModel(PreTrainedModel):
     config_class = LanceAIConfig
     base_model_prefix = "model"
@@ -260,6 +268,7 @@ class LanceAIPreTrainedModel(PreTrainedModel):
     _supports_flex_attn = True
     _can_compile_fullgraph = True
 
+
 class LanceAIModel(LanceAIPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
@@ -335,6 +344,7 @@ class LanceAIModel(LanceAIPreTrainedModel):
         mask = torch.triu(mask, diagonal=1 + past_len)
         return mask[None, None, :, :]
 
+
 class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
     _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}
 
@@ -404,6 +414,7 @@ class LanceAI(LanceAIPreTrainedModel, GenerationMixin):
             reordered.append((layer_k.index_select(0, beam_idx), layer_v.index_select(0, beam_idx)))
         return reordered
 
+
 CONFIG_MAPPING.register("lance_ai", LanceAIConfig)
 MODEL_FOR_CAUSAL_LM_MAPPING.register(LanceAIConfig, LanceAI)
 LanceAIConfig.register_for_auto_class("AutoConfig")