Super-squash branch 'main' using huggingface_hub

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,203 @@
+---
+library_name: transformers
+license: mit
+datasets:
+- EleutherAI/SmolLM2-1.7B-stage-4-100B
+language:
+- en
+---
+
+# NeoBERT Model
+
+This is a NeoBERT model trained with [pszemraj/NeoBERT](https://github.com/pszemraj/NeoBERT) and exported to `transformers` format.
+
+## Model Details
+- **Architecture**: NeoBERT
+- **Hidden Size**: 768
+- **Layers**: 12
+- **Attention Heads**: 12
+- **Vocab Size**: 31999
+- **Max Length**: 4096
+- **Dtype**: float32
+
+## Usage
+
+### For Masked Language Modeling (Fill-Mask)
+
+```python
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+import torch
+
+repo_id = "BEE-spoke-data/neobert-100k-test"
+tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModelForMaskedLM.from_pretrained(repo_id, trust_remote_code=True)
+
+# Example: Fill in masked tokens
+text = "NeoBERT is the most [MASK] model of its kind!"
+
+# Tokenize (handling Metaspace tokenizer's space tokens)
+inputs = tokenizer(text, return_tensors="pt")
+input_ids = inputs["input_ids"][0].tolist()
+
+# Remove extra space tokens before [MASK] if present (Metaspace tokenizer quirk)
+cleaned_ids = []
+for i, token_id in enumerate(input_ids):
+    if token_id == 454 and i < len(input_ids) - 1 and input_ids[i + 1] == tokenizer.mask_token_id:
+        continue
+    cleaned_ids.append(token_id)
+
+if len(cleaned_ids) != len(input_ids):
+    inputs["input_ids"] = torch.tensor([cleaned_ids])
+    inputs["attention_mask"] = torch.ones_like(inputs["input_ids"])
+
+# Get predictions
+with torch.no_grad():
+    outputs = model(**inputs)
+    mask_pos = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(as_tuple=True)[1][0]
+    predictions = outputs.logits[0, mask_pos].topk(5)
+
+# Display top predictions
+for idx, score in zip(predictions.indices, predictions.values):
+    token = tokenizer.decode([idx])
+    print(f"{token}: {score:.2f}")
+```
+
+### For Embeddings / Feature Extraction
+
+```python
+from transformers import AutoModel, AutoTokenizer
+
+repo_id = "BEE-spoke-data/neobert-100k-test"
+tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(repo_id, trust_remote_code=True)
+
+# Example: Generate embeddings
+text = "NeoBERT is an efficient transformer model!"
+inputs = tokenizer(text, return_tensors="pt")
+outputs = model(**inputs)
+
+# Get CLS token embedding
+cls_embedding = outputs.last_hidden_state[:, 0, :]
+print(f"Embedding shape: {cls_embedding.shape}")
+```
+
+## Training Configuration
+
+<details>
+  <summary><strong>Full Config</strong> (click to expand)</summary>
+
+Full training config:
+
+```yaml
+model:
+  hidden_size: 768
+  num_hidden_layers: 12
+  num_attention_heads: 12
+  intermediate_size: 3072
+  max_position_embeddings: 4096
+  vocab_size: 31999
+  rope: true
+  rms_norm: true
+  hidden_act: swiglu
+  dropout_prob: 0.05
+  norm_eps: 1.0e-05
+  embedding_init_range: 0.02
+  decoder_init_range: 0.02
+  classifier_init_range: 0.02
+  flash_attention: true
+  ngpt: false
+  base_scale: 0.03227486121839514
+  pad_token_id: 0
+dataset:
+  name: EleutherAI/SmolLM2-1.7B-stage-4-100B
+  path: ''
+  num_workers: 4
+  streaming: true
+  cache_dir: null
+  max_seq_length: 1024
+  validation_split: null
+  train_split: train
+  eval_split: train[:1%]
+  num_proc: 8
+  shuffle_buffer_size: 10000
+  pre_tokenize: false
+  pre_tokenize_output: null
+  load_all_from_disk: false
+  force_redownload: false
+  pretraining_prob: 0.3
+  min_length: 512
+tokenizer:
+  name: BEE-spoke-data/wordpiece-tokenizer-32k-en_code-msp
+  path: null
+  max_length: 1024
+  padding: max_length
+  truncation: true
+  vocab_size: 31999
+optimizer:
+  name: adamw
+  lr: 0.0001
+  weight_decay: 0.01
+  betas:
+  - 0.9
+  - 0.98
+  eps: 1.0e-08
+scheduler:
+  name: cosine
+  warmup_steps: 5000
+  total_steps: null
+  num_cycles: 0.5
+  decay_steps: 50000
+  warmup_percent: null
+  decay_percent: null
+trainer:
+  per_device_train_batch_size: 16
+  per_device_eval_batch_size: 16
+  gradient_accumulation_steps: 4
+  max_steps: 100000
+  save_steps: 10000
+  eval_steps: 5000
+  logging_steps: 25
+  output_dir: ./outputs/neobert_100m_100k
+  overwrite_output_dir: true
+  bf16: true
+  gradient_checkpointing: false
+  gradient_clipping: null
+  mixed_precision: 'no'
+  seed: 42
+  resume_from_checkpoint: false
+  disable_tqdm: false
+  dataloader_num_workers: 0
+  use_cpu: false
+  report_to:
+  - wandb
+  tf32: true
+  max_ckpt: 3
+  train_batch_size: 16
+  eval_batch_size: 32
+datacollator:
+  mlm_probability: 0.2
+  pad_to_multiple_of: 8
+wandb:
+  project: neobert-pretraining
+  entity: null
+  name: neobert-100m-100k
+  tags: []
+  mode: online
+  log_interval: 100
+  resume: never
+  dir: logs/wandb
+task: pretraining
+accelerate_config_file: null
+mixed_precision: bf16
+mteb_task_type: all
+mteb_batch_size: 32
+mteb_pooling: mean
+mteb_overwrite_results: false
+pretrained_checkpoint: latest
+use_deepspeed: true
+seed: 69
+debug: false
+
+```
+
+</details>

config.json

@@ -0,0 +1,24 @@
+{
+  "architectures": [
+    "NeoBERTLMHead"
+  ],
+  "model_type": "neobert",
+  "auto_map": {
+    "AutoConfig": "model.NeoBERTConfig",
+    "AutoModel": "model.NeoBERT",
+    "AutoModelForMaskedLM": "model.NeoBERTLMHead",
+    "AutoModelForSequenceClassification": "model.NeoBERTForSequenceClassification"
+  },
+  "hidden_size": 768,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 12,
+  "intermediate_size": 3072,
+  "vocab_size": 31999,
+  "max_length": 4096,
+  "embedding_init_range": 0.02,
+  "decoder_init_range": 0.02,
+  "norm_eps": 1e-05,
+  "pad_token_id": 0,
+  "torch_dtype": "float32",
+  "transformers_version": "4.55.0"
+}

model.py

@@ -0,0 +1,505 @@
+# From https://github.com/facebookresearch/llama/blob/main/llama/model.py
+
+from typing import Optional
+
+import numpy as np
+import torch
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.nn.functional import scaled_dot_product_attention
+
+try:
+    from xformers.ops import SwiGLU
+
+    XFORMERS_AVAILABLE = True
+except ImportError:
+    XFORMERS_AVAILABLE = False
+
+try:
+    from flash_attn.flash_attn_interface import flash_attn_varlen_func
+
+    FLASH_ATTN_AVAILABLE = True
+except ImportError:
+    FLASH_ATTN_AVAILABLE = False
+
+from transformers import (
+    DataCollatorForLanguageModeling,
+    PretrainedConfig,
+    PreTrainedModel,
+)
+from transformers.modeling_outputs import (
+    BaseModelOutput,
+    MaskedLMOutput,
+    SequenceClassifierOutput,
+)
+
+from .rotary import apply_rotary_emb, precompute_freqs_cis
+
+
+class DataCollatorWithPacking(DataCollatorForLanguageModeling):
+    def __init__(self, pack_sequences=False, **kwargs):
+        super().__init__(**kwargs)
+        self.pack_sequences = pack_sequences
+
+    def __call__(self, batch):
+        if self.pack_sequences:
+            # Add position_ids if not present
+            if "position_ids" not in batch[0]:
+                for item in batch:
+                    item["position_ids"] = list(range(len(item["input_ids"])))
+
+            # Pack the sequences into a single list
+            input_ids_list = [item["input_ids"] for item in batch]
+            position_ids_list = [item["position_ids"] for item in batch]
+            seqlens = np.array([0] + [len(ids) for ids in input_ids_list])
+
+            packed_batch = {
+                "position_ids": np.concatenate(position_ids_list, axis=0),
+                "input_ids": np.concatenate(input_ids_list, axis=0),
+                "cu_seqlens": np.cumsum(seqlens),
+                "max_seqlen": max(seqlens),
+            }
+
+            batch = super().__call__([packed_batch])
+            batch["cu_seqlens"] = batch["cu_seqlens"].to(torch.int32).squeeze()
+        else:
+            batch = super().__call__(batch)
+            batch["attention_mask"] = batch["attention_mask"].to(torch.bool)
+
+        return batch
+
+
+class NeoBERTConfig(PretrainedConfig):
+    model_type = "neobert"
+
+    # All config parameters must have a default value.
+    def __init__(
+        self,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 28,
+        num_attention_heads: int = 12,
+        intermediate_size: int = 3072,
+        embedding_init_range: float = 0.02,
+        decoder_init_range: float = 0.02,
+        norm_eps: float = 1e-06,
+        vocab_size: int = 30522,
+        pad_token_id: int = 0,
+        max_length: int = 1024,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        if hidden_size % num_attention_heads != 0:
+            raise ValueError("Hidden size must be divisible by the number of heads.")
+        self.dim_head = hidden_size // num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.embedding_init_range = embedding_init_range
+        self.decoder_init_range = decoder_init_range
+        self.norm_eps = norm_eps
+        self.vocab_size = vocab_size
+        self.pad_token_id = pad_token_id
+        self.max_length = max_length
+        self.kwargs = kwargs
+
+
+# Adapted from transformers.models.llama.modeling_llama.LlamaMLP
+class NeobertMLP(nn.Module):
+    def __init__(self, hidden_size, intermediate_size, bias=False):
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.w12 = nn.Linear(self.hidden_size, 2 * self.intermediate_size, bias=bias)
+        self.w3 = nn.Linear(self.intermediate_size, self.hidden_size, bias=bias)
+        self.act_fn = nn.SiLU()
+
+    def forward(self, x):
+        w1, w2 = self.w12(x).chunk(2, dim=-1)
+        w3 = self.w3(self.act_fn(w1) * w2)
+        return w3
+
+
+class EncoderBlock(nn.Module):
+    """Transformer encoder block."""
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__()
+
+        self.config = config
+
+        # Attention
+        self.qkv = nn.Linear(
+            in_features=config.hidden_size,
+            out_features=config.hidden_size * 3,
+            bias=False,
+        )
+        self.wo = nn.Linear(
+            in_features=config.hidden_size, out_features=config.hidden_size, bias=False
+        )
+
+        # Feedforward network
+        multiple_of = 8
+        intermediate_size = int(2 * config.intermediate_size / 3)
+        intermediate_size = multiple_of * (
+            (intermediate_size + multiple_of - 1) // multiple_of
+        )
+        if XFORMERS_AVAILABLE:
+            self.ffn = SwiGLU(
+                config.hidden_size, intermediate_size, config.hidden_size, bias=False
+            )
+        else:
+            self.ffn = NeobertMLP(config.hidden_size, intermediate_size, bias=False)
+
+        # Layer norms
+        self.attention_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
+        self.ffn_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        output_attentions: bool,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+    ):
+        # Attention
+        attn_output, attn_weights = self._att_block(
+            self.attention_norm(x),
+            attention_mask,
+            freqs_cis,
+            output_attentions,
+            max_seqlen,
+            cu_seqlens,
+        )
+
+        # Residual
+        x = x + attn_output
+
+        # Feed-forward
+        x = x + self.ffn(self.ffn_norm(x))
+
+        return x, attn_weights
+
+    def _att_block(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor,
+        freqs_cis: torch.Tensor,
+        output_attentions: bool,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+    ):
+        batch_size, seq_len, _ = x.shape
+
+        xq, xk, xv = (
+            self.qkv(x)
+            .view(
+                batch_size,
+                seq_len,
+                self.config.num_attention_heads,
+                self.config.dim_head * 3,
+            )
+            .chunk(3, axis=-1)
+        )
+
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis)
+
+        # Attn block
+        attn_weights = None
+
+        # Flash attention if the tensors are packed
+        if cu_seqlens is not None:
+            attn = flash_attn_varlen_func(
+                q=xq.squeeze(0),
+                k=xk.squeeze(0),
+                v=xv.squeeze(0),
+                cu_seqlens_q=cu_seqlens,
+                cu_seqlens_k=cu_seqlens,
+                max_seqlen_q=max_seqlen,
+                max_seqlen_k=max_seqlen,
+                dropout_p=0.0,
+                causal=False,
+            )
+        # Eager attention if attention weights are needed in the output
+        elif output_attentions:
+            attn_weights = (
+                xq.permute(0, 2, 1, 3) @ xk.permute(0, 2, 3, 1) / (xq.size(-1) ** 0.5)
+            )
+            if attention_mask is not None:
+                attn_weights = attn_weights * attention_mask
+            attn_weights = attn_weights.softmax(-1)
+            attn = attn_weights @ xv.permute(0, 2, 1, 3)
+            attn = attn.transpose(1, 2)
+        # Fall back to SDPA otherwise
+        else:
+            attn = scaled_dot_product_attention(
+                query=xq.transpose(1, 2),
+                key=xk.transpose(1, 2),
+                value=xv.transpose(1, 2),
+                attn_mask=attention_mask.bool() if attention_mask is not None else None,
+                dropout_p=0,
+            ).transpose(1, 2)
+
+        return (
+            self.wo(
+                attn.reshape(
+                    batch_size,
+                    seq_len,
+                    self.config.num_attention_heads * self.config.dim_head,
+                )
+            ),
+            attn_weights,
+        )
+
+
+class NeoBERTPreTrainedModel(PreTrainedModel):
+    config_class = NeoBERTConfig
+    base_model_prefix = "model"
+    _supports_cache_class = True
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.uniform_(
+                -self.config.decoder_init_range, self.config.decoder_init_range
+            )
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.uniform_(
+                -self.config.embedding_init_range, self.config.embedding_init_range
+            )
+
+
+class NeoBERT(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.encoder = nn.Embedding(
+            config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id
+        )
+
+        # Ensures freqs_cis is moved to the same devices as the model. Non-persistent buffers are not saved in the state_dict.
+        freqs_cis = precompute_freqs_cis(
+            config.hidden_size // config.num_attention_heads, config.max_length
+        )
+        self.register_buffer("freqs_cis", freqs_cis, persistent=False)
+
+        self.transformer_encoder = nn.ModuleList()
+        for _ in range(config.num_hidden_layers):
+            self.transformer_encoder.append(EncoderBlock(config))
+
+        self.layer_norm = nn.RMSNorm(config.hidden_size, config.norm_eps)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        **kwargs,
+    ):
+        # Initialize
+        hidden_states, attentions = [], []
+
+        # Expand and repeat: (Batch, Length) -> (Batch, Heads, Length, Length)
+        if attention_mask is not None:
+            attention_mask = (
+                attention_mask.unsqueeze(1)
+                .unsqueeze(1)
+                .repeat(1, self.config.num_attention_heads, attention_mask.size(-1), 1)
+            )
+
+        # Checks to be done if inputs are packed sequences
+        if cu_seqlens is not None:
+            assert FLASH_ATTN_AVAILABLE, (
+                "Flash-attention is not available. Please ''pip install flash_attn'', or provide un-packed sequences."
+            )
+            assert not output_attentions, (
+                "Output attentions is not supported when sequences are packed."
+            )
+            assert max_seqlen is not None, (
+                "Missing max_seqlen. It must be provided when cu_seqlens are not None."
+            )
+            assert input_ids.shape[0] == 1, (
+                "Cumulative sequence lengths are provided but input_ids are not packed."
+            )
+            assert input_ids.is_cuda, (
+                "Packing uses an implementation of flash-attention and is only supported on GPU."
+            )
+
+        # RoPE
+        freqs_cis = (
+            self.freqs_cis[position_ids]
+            if position_ids is not None
+            else self.freqs_cis[: input_ids.shape[1]].unsqueeze(0)
+        )
+
+        # Embedding
+        x = self.encoder(input_ids)
+
+        # Transformer encoder
+        for layer in self.transformer_encoder:
+            x, attn = layer(
+                x, attention_mask, freqs_cis, output_attentions, max_seqlen, cu_seqlens
+            )
+            if output_hidden_states:
+                hidden_states.append(x)
+            if output_attentions:
+                attentions.append(attn)
+
+        # Final normalization layer
+        x = self.layer_norm(x)
+
+        # Return the output of the last hidden layer
+        return BaseModelOutput(
+            last_hidden_state=x,
+            hidden_states=hidden_states if output_hidden_states else None,
+            attentions=attentions if output_attentions else None,
+        )
+
+
+class NeoBERTLMHead(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.model = NeoBERT(config)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)
+
+        self.post_init()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        **kwargs,
+    ):
+        output = self.model.forward(
+            input_ids,
+            position_ids,
+            max_seqlen,
+            cu_seqlens,
+            attention_mask,
+            output_hidden_states,
+            output_attentions,
+        )
+        logits = self.decoder(output.last_hidden_state)
+
+        return MaskedLMOutput(
+            hidden_states=output.hidden_states if output_hidden_states else None,
+            attentions=output.attentions if output_attentions else None,
+            logits=logits,
+        )
+
+
+class NeoBERTForSequenceClassification(NeoBERTPreTrainedModel):
+    config_class = NeoBERTConfig
+
+    def __init__(self, config: NeoBERTConfig):
+        super().__init__(config)
+
+        self.config = config
+
+        self.num_labels = getattr(config, "num_labels", 2)
+        self.classifier_dropout = getattr(config, "classifier_dropout", 0.1)
+        self.classifier_init_range = getattr(config, "classifier_init_range", 0.02)
+
+        self.model = NeoBERT(config)
+
+        self.dense = nn.Linear(self.config.hidden_size, self.config.hidden_size)
+        self.dropout = nn.Dropout(self.classifier_dropout)
+        self.classifier = nn.Linear(self.config.hidden_size, self.num_labels)
+
+        self.post_init()
+
+    def _init_weights(self, module):
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=self.classifier_init_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        position_ids: torch.Tensor = None,
+        max_seqlen: int = None,
+        cu_seqlens: torch.Tensor = None,
+        attention_mask: torch.Tensor = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        labels: Optional[torch.Tensor] = None,
+        return_dict: Optional[bool] = None,
+    ):
+        output = self.model.forward(
+            input_ids,
+            position_ids,
+            max_seqlen,
+            cu_seqlens,
+            attention_mask,
+            output_hidden_states,
+            output_attentions,
+        )
+        hidden_states = output.last_hidden_state
+
+        x = hidden_states[:, 0, :]
+        x = self.dropout(x)
+        x = self.dense(x)
+        x = torch.tanh(x)
+        x = self.dropout(x)
+
+        logits = self.classifier(x)
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(logits, labels)
+
+        if not return_dict:
+            result = (logits,)
+            return ((loss,) + result) if loss is not None else result
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states if output_hidden_states else None,
+            attentions=output.attentions if output_attentions else None,
+        )

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:33b6bf91a447ce603db6ae080de1a67dde749fd867f471029a22a03cbdfa9eab
+size 536553948

pytorch_model.bin

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

rotary.py

@@ -0,0 +1,62 @@
+# From https://github.com/facebookresearch/llama/blob/main/llama/model.py
+
+from typing import Tuple
+
+import torch
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    """
+    Precompute the frequency tensor for complex exponentials (cis) with given dimensions.
+
+    This function calculates a frequency tensor with complex exponentials using the given dimension 'dim'
+    and the end index 'end'. The 'theta' parameter scales the frequencies.
+    The returned tensor contains complex values in complex64 data type.
+
+    Args:
+        dim (int): Dimension of the frequency tensor.
+        end (int): End index for precomputing frequencies.
+        theta (float, optional): Scaling factor for frequency computation. Defaults to 10000.0.
+
+    Returns:
+        torch.Tensor: Precomputed frequency tensor with complex exponentials.
+    """
+
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)
+    freqs = torch.outer(t, freqs).float()
+    return torch.polar(torch.ones_like(freqs), freqs)
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    assert freqs_cis.shape[1:] == (x.shape[1], x.shape[-1])
+    return freqs_cis.contiguous().unsqueeze(2)
+
+
+def apply_rotary_emb(
+    xq: torch.Tensor,
+    xk: torch.Tensor,
+    freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Apply rotary embeddings to input tensors using the given frequency tensor.
+
+    This function applies rotary embeddings to the given query 'xq' and key 'xk' tensors using the provided
+    frequency tensor 'freqs_cis'. The input tensors are reshaped as complex numbers, and the frequency tensor
+    is reshaped for broadcasting compatibility. The resulting tensors contain rotary embeddings and are
+    returned as real tensors.
+
+    Args:
+        xq (torch.Tensor): Query tensor to apply rotary embeddings.
+        xk (torch.Tensor): Key tensor to apply rotary embeddings.
+        freqs_cis (torch.Tensor): Precomputed frequency tensor for complex exponentials.
+
+    Returns:
+        Tuple[torch.Tensor, torch.Tensor]: Tuple of modified query tensor and key tensor with rotary embeddings.
+    """
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)

special_tokens_map.json

@@ -0,0 +1,51 @@
+{
+  "bos_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "cls_token": {
+    "content": "[CLS]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "mask_token": {
+    "content": "[MASK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "[PAD]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "sep_token": {
+    "content": "[SEP]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "[UNK]",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,60 @@
+{
+  "added_tokens_decoder": {
+    "0": {
+      "content": "[UNK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "[CLS]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "[SEP]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "[PAD]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "[MASK]",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "bos_token": "[CLS]",
+  "clean_up_tokenization_spaces": true,
+  "cls_token": "[CLS]",
+  "eos_token": "[SEP]",
+  "extra_special_tokens": {},
+  "mask_token": "[MASK]",
+  "max_length": 1024,
+  "model_max_length": 4096,
+  "pad_token": "[PAD]",
+  "sep_token": "[SEP]",
+  "stride": 0,
+  "tokenizer_class": "PreTrainedTokenizerFast",
+  "truncation_side": "right",
+  "truncation_strategy": "longest_first",
+  "unk_token": "[UNK]",
+  "vocab_size": 31999
+}