Upload CubeLM

CubeConfig.py

@@ -0,0 +1,33 @@
+#from transformers import PretrainedConfig
+from transformers import GPT2Config
+from cubeLM.tokenizer import vocab
+
+vocab_size = len(vocab)
+
+
+class CubeConfig(GPT2Config):
+    model_type = "CubeLM"
+
+    def __init__(
+        self,
+        vocab_size=vocab_size,
+        bos_token_id=vocab_size - 1,
+        eos_token_id=vocab_size - 1,
+        pad_token_id=vocab_size - 1,
+        n_positions=40,
+        n_embd=512,
+        n_layer=8,
+        n_head=8,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.n_embd = n_embd
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.pad_token_id = pad_token_id
+
+

CubeLM.py

@@ -0,0 +1,119 @@
+from dataclasses import dataclass
+from typing import Optional
+import torch
+import torch.nn as nn
+from cubeLM.CubeConfig import CubeConfig
+from transformers import (
+    GPT2Model,
+    GenerationMixin,
+    GPT2PreTrainedModel,
+    PreTrainedModel
+)
+from transformers.utils import ModelOutput
+from train_scripts.utils import IGNORE_INDEX
+
+
+@dataclass
+class CubeLMOutput(ModelOutput):
+    total_loss: Optional[torch.FloatTensor] = None
+    lm_loss: Optional[torch.FloatTensor] = None
+    cube_loss: Optional[torch.FloatTensor] = None
+    logits: Optional[torch.FloatTensor] = None
+    cube_logits: Optional[torch.FloatTensor] = None
+
+
+class CubeLM(PreTrainedModel):
+    config_class = CubeConfig
+    _no_split_modules = ["GPT2Block"]
+
+    def __init__(self, config, task="sft", num_heads=24, num_classes=6):
+        super().__init__(config)
+        self.transformer = GPT2Model(config)
+        assert task in ["sft", "pretrain", "joint"]
+
+        self.task = task
+        self.alpha = None
+        if hasattr(config, "alpha"):
+            self.alpha = config.alpha
+        self.vocab_size = config.vocab_size
+        self.lm_head = None
+        self.cube_heads = None
+        if task in ["sft", "joint"]:
+            self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        if task in ["pretrain", "joint"]:
+            self.cube_heads = nn.Linear(
+                config.n_embd, num_heads * num_classes, bias=False
+            )
+            self.num_heads = num_heads
+            self.num_classes = num_classes
+        self.config = config
+
+    @classmethod
+    def can_generate(cls):
+        return True
+
+    def forward(
+        self, input_ids, attention_mask=None, labels=None, cube_states=None, **kwargs
+    ):
+        outputs = self.transformer(
+            input_ids,
+            attention_mask=attention_mask,
+        )
+        # [batch, seq, d_model]
+        hidden_states = outputs.last_hidden_state
+
+        lm_logits = None
+        lm_loss = None
+        # [batch, seq, n_heads * n_classes]
+        if self.lm_head:
+            lm_logits = self.lm_head(hidden_states)
+
+            if labels is not None:
+                shift_logits = lm_logits[:, :-1, :].contiguous()
+                shift_labels = input_ids[:, 1:].contiguous()
+                loss_fn = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)
+                lm_loss = loss_fn(
+                    shift_logits.view(-1, self.vocab_size),
+                    shift_labels.view(-1),
+                )
+
+        cube_logits = None
+        cube_loss = None
+        if self.cube_heads:
+            cube_logits = self.cube_heads(hidden_states)
+
+            if cube_states is not None:
+                cube_logits = cube_logits.view(
+                    hidden_states.size(0),
+                    hidden_states.size(1),
+                    self.num_heads,
+                    self.num_classes,
+                )
+                # [batch * seq, n_heads, n_classes]
+                _logits = cube_logits.view(-1, self.num_heads, self.num_classes)
+                # [batch * seq, n_heads]
+                _labels = cube_states.view(-1, self.num_heads)
+                loss_fn = nn.CrossEntropyLoss(ignore_index=IGNORE_INDEX)
+                losses = []
+                for head_idx in range(self.num_heads):
+                    losses.append(
+                        loss_fn(_logits[:, head_idx, :], _labels[:, head_idx])
+                    )
+                cube_loss = sum(losses) / self.num_heads
+
+        total_loss = None
+        if lm_loss is not None and cube_loss is not None:
+            assert self.alpha is not None
+            total_loss = lm_loss + self.alpha * cube_loss
+        elif lm_loss is not None:
+            total_loss = lm_loss
+        elif cube_loss is not None:
+            total_loss = cube_loss
+
+        return CubeLMOutput(
+            total_loss=total_loss,
+            lm_loss=lm_loss,
+            cube_loss=cube_loss,
+            logits=lm_logits,
+            cube_logits=cube_logits,
+        )

README.md

@@ -0,0 +1,199 @@
+---
+library_name: transformers
+tags: []
+---
+
+# Model Card for Model ID
+
+<!-- Provide a quick summary of what the model is/does. -->
+
+
+
+## Model Details
+
+### Model Description
+
+<!-- Provide a longer summary of what this model is. -->
+
+This is the model card of a 🤗 transformers model that has been pushed on the Hub. This model card has been automatically generated.
+
+- **Developed by:** [More Information Needed]
+- **Funded by [optional]:** [More Information Needed]
+- **Shared by [optional]:** [More Information Needed]
+- **Model type:** [More Information Needed]
+- **Language(s) (NLP):** [More Information Needed]
+- **License:** [More Information Needed]
+- **Finetuned from model [optional]:** [More Information Needed]
+
+### Model Sources [optional]
+
+<!-- Provide the basic links for the model. -->
+
+- **Repository:** [More Information Needed]
+- **Paper [optional]:** [More Information Needed]
+- **Demo [optional]:** [More Information Needed]
+
+## Uses
+
+<!-- Address questions around how the model is intended to be used, including the foreseeable users of the model and those affected by the model. -->
+
+### Direct Use
+
+<!-- This section is for the model use without fine-tuning or plugging into a larger ecosystem/app. -->
+
+[More Information Needed]
+
+### Downstream Use [optional]
+
+<!-- This section is for the model use when fine-tuned for a task, or when plugged into a larger ecosystem/app -->
+
+[More Information Needed]
+
+### Out-of-Scope Use
+
+<!-- This section addresses misuse, malicious use, and uses that the model will not work well for. -->
+
+[More Information Needed]
+
+## Bias, Risks, and Limitations
+
+<!-- This section is meant to convey both technical and sociotechnical limitations. -->
+
+[More Information Needed]
+
+### Recommendations
+
+<!-- This section is meant to convey recommendations with respect to the bias, risk, and technical limitations. -->
+
+Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
+
+## How to Get Started with the Model
+
+Use the code below to get started with the model.
+
+[More Information Needed]
+
+## Training Details
+
+### Training Data
+
+<!-- This should link to a Dataset Card, perhaps with a short stub of information on what the training data is all about as well as documentation related to data pre-processing or additional filtering. -->
+
+[More Information Needed]
+
+### Training Procedure
+
+<!-- This relates heavily to the Technical Specifications. Content here should link to that section when it is relevant to the training procedure. -->
+
+#### Preprocessing [optional]
+
+[More Information Needed]
+
+
+#### Training Hyperparameters
+
+- **Training regime:** [More Information Needed] <!--fp32, fp16 mixed precision, bf16 mixed precision, bf16 non-mixed precision, fp16 non-mixed precision, fp8 mixed precision -->
+
+#### Speeds, Sizes, Times [optional]
+
+<!-- This section provides information about throughput, start/end time, checkpoint size if relevant, etc. -->
+
+[More Information Needed]
+
+## Evaluation
+
+<!-- This section describes the evaluation protocols and provides the results. -->
+
+### Testing Data, Factors & Metrics
+
+#### Testing Data
+
+<!-- This should link to a Dataset Card if possible. -->
+
+[More Information Needed]
+
+#### Factors
+
+<!-- These are the things the evaluation is disaggregating by, e.g., subpopulations or domains. -->
+
+[More Information Needed]
+
+#### Metrics
+
+<!-- These are the evaluation metrics being used, ideally with a description of why. -->
+
+[More Information Needed]
+
+### Results
+
+[More Information Needed]
+
+#### Summary
+
+
+
+## Model Examination [optional]
+
+<!-- Relevant interpretability work for the model goes here -->
+
+[More Information Needed]
+
+## Environmental Impact
+
+<!-- Total emissions (in grams of CO2eq) and additional considerations, such as electricity usage, go here. Edit the suggested text below accordingly -->
+
+Carbon emissions can be estimated using the [Machine Learning Impact calculator](https://mlco2.github.io/impact#compute) presented in [Lacoste et al. (2019)](https://arxiv.org/abs/1910.09700).
+
+- **Hardware Type:** [More Information Needed]
+- **Hours used:** [More Information Needed]
+- **Cloud Provider:** [More Information Needed]
+- **Compute Region:** [More Information Needed]
+- **Carbon Emitted:** [More Information Needed]
+
+## Technical Specifications [optional]
+
+### Model Architecture and Objective
+
+[More Information Needed]
+
+### Compute Infrastructure
+
+[More Information Needed]
+
+#### Hardware
+
+[More Information Needed]
+
+#### Software
+
+[More Information Needed]
+
+## Citation [optional]
+
+<!-- If there is a paper or blog post introducing the model, the APA and Bibtex information for that should go in this section. -->
+
+**BibTeX:**
+
+[More Information Needed]
+
+**APA:**
+
+[More Information Needed]
+
+## Glossary [optional]
+
+<!-- If relevant, include terms and calculations in this section that can help readers understand the model or model card. -->
+
+[More Information Needed]
+
+## More Information [optional]
+
+[More Information Needed]
+
+## Model Card Authors [optional]
+
+[More Information Needed]
+
+## Model Card Contact
+
+[More Information Needed]

config.json

@@ -0,0 +1,36 @@
+{
+  "activation_function": "gelu_new",
+  "architectures": [
+    "CubeLM"
+  ],
+  "attn_pdrop": 0.1,
+  "auto_map": {
+    "AutoConfig": "CubeConfig.CubeConfig",
+    "AutoModelForCausalLM": "CubeLM.CubeLM"
+  },
+  "bos_token_id": 15,
+  "embd_pdrop": 0.1,
+  "eos_token_id": 15,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "model_type": "CubeLM",
+  "n_embd": 512,
+  "n_head": 8,
+  "n_inner": null,
+  "n_layer": 8,
+  "n_positions": 50,
+  "pad_token_id": 15,
+  "reorder_and_upcast_attn": false,
+  "resid_pdrop": 0.1,
+  "scale_attn_by_inverse_layer_idx": false,
+  "scale_attn_weights": true,
+  "summary_activation": null,
+  "summary_first_dropout": 0.1,
+  "summary_proj_to_labels": true,
+  "summary_type": "cls_index",
+  "summary_use_proj": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.49.0",
+  "use_cache": true,
+  "vocab_size": 16
+}

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 15,
+  "eos_token_id": 15,
+  "pad_token_id": 15,
+  "transformers_version": "4.49.0"
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72c0fada7f0ef9dc0a4baa0a7eaf10dceb1bb5b044e6a825d225d480f49117eb
+size 101058360