Upload LanceAI

config.json

@@ -2,11 +2,15 @@
   "architectures": [
     "LanceAI"
   ],
+  "auto_map": {
+    "AutoConfig": "lance_ai_model.LanceAIConfig",
+    "AutoModelForCausalLM": "lance_ai_model.LanceAI"
+  },
   "hidden_size": 2048,
   "model_type": "lance_ai",
   "num_heads": 16,
   "num_layers": 24,
   "torch_dtype": "float32",
-  "transformers_version": "4.50.3",
+  "transformers_version": "4.51.3",
   "vocab_size": 100277
 }

generation_config.json

@@ -1,4 +1,4 @@
 {
   "_from_model_config": true,
-  "transformers_version": "4.50.3"
+  "transformers_version": "4.51.3"
 }

lance_ai_model.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin
+from transformers.modeling_outputs import CausalLMOutputWithPast
+from transformers.models.auto.configuration_auto import CONFIG_MAPPING
+from transformers.models.auto.modeling_auto import MODEL_FOR_CAUSAL_LM_MAPPING
+
+class LanceAIConfig(PretrainedConfig):
+    model_type = "lance_ai"
+    def __init__(self, vocab_size=50257, hidden_size=2048, num_layers=24, num_heads=16, architectures=["LanceAI"], **kwargs):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.num_heads = num_heads
+        self.architectures = architectures
+
+class LanceAI(PreTrainedModel, GenerationMixin):
+    config_class = LanceAIConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+
+        self.encoder = nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(d_model=config.hidden_size, nhead=config.num_heads),
+            num_layers=config.num_layers
+        )
+        self.decoder = nn.TransformerDecoder(
+            nn.TransformerDecoderLayer(d_model=config.hidden_size, nhead=config.num_heads),
+            num_layers=config.num_layers
+        )
+
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size)
+        self.loss_fct = nn.CrossEntropyLoss(ignore_index=-100)
+
+        self.init_weights()
+
+    def forward(self, input_ids=None, labels=None, inputs_embeds=None, return_dict=True):
+        embeddings = self.embedding(input_ids) if inputs_embeds is None else inputs_embeds
+        encoder_output = self.encoder(embeddings)
+        decoder_output = self.decoder(embeddings, encoder_output)
+
+        logits = self.lm_head(decoder_output)
+        loss = None
+
+        if labels is not None:
+            shift_logits = logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            shift_labels = torch.clamp(shift_labels, max=self.config.vocab_size - 1)
+            loss = self.loss_fct(shift_logits.view(-1, self.config.vocab_size), shift_labels.view(-1))
+
+        if return_dict:
+            return CausalLMOutputWithPast(loss=loss, logits=logits)
+
+        return (loss, logits) if loss is not None else logits
+
+# Register with Hugging Face
+CONFIG_MAPPING.register("lance_ai", LanceAIConfig)
+MODEL_FOR_CAUSAL_LM_MAPPING.register(LanceAIConfig, LanceAI)
+LanceAIConfig.register_for_auto_class("AutoConfig")
+LanceAI.register_for_auto_class("AutoModelForCausalLM")