Upload model

config.json

@@ -0,0 +1,14 @@
+{
+  "_name_or_path": "./checkpoints/mcqa_model_downsampled_bf16",
+  "architectures": [
+    "MCQHFModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "model.MyCustomConfig",
+    "AutoModelForCausalLM": "model.MCQHFModel"
+  },
+  "model_type": "mcq_hf_model",
+  "name_model": "apple/OpenELM-450M-Instruct",
+  "torch_dtype": "float32",
+  "transformers_version": "4.40.2"
+}

model.py

@@ -0,0 +1,95 @@
+import torch
+import torch.nn as nn
+from transformers import (
+    PreTrainedModel,
+    PretrainedConfig,
+    AutoConfig,
+    AutoModel,
+    AutoModelForCausalLM,
+    AutoTokenizer,
+)
+
+import pandas as pd
+from datasets import load_dataset, Dataset, DatasetDict
+
+
+class MCQModel(nn.Module):
+    def __init__(self, name_model):
+        super(MCQModel, self).__init__()
+        self.model = AutoModelForCausalLM.from_pretrained(
+            name_model,
+            trust_remote_code=True,
+            output_hidden_states=True,
+        )
+        self.classifier = nn.Linear(
+            self.model.config.model_dim, 4
+        )  # 4 classes for 'A', 'B', 'C', 'D'
+
+    def forward(self, input_ids, attention_mask=None, labels=None,
+        position_ids = None,
+        past_key_values = None,
+        inputs_embeds = None,
+        use_cache  = None,
+        output_attentions  = None,
+        output_hidden_states = None,
+        return_dict  = None,
+        cache_position = None):
+
+        print("INSIDE CUSTOM MODEL LABELS: ", labels)
+        outputs = self.model(input_ids, attention_mask=attention_mask, position_ids= position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds,
+                              output_attentions=output_attentions, cache_position=cache_position, use_cache=use_cache, return_dict=return_dict, output_hidden_states=output_hidden_states)
+        print("OUTPUT KEYS: " , outputs.keys())
+        # print(outputs.logits.shape)
+        # # print(outputs.hidden_states)
+        # print(outputs.hidden_states[0].shape)
+        # print(outputs.hidden_states[1].shape)
+        # print(len(outputs.hidden_states))
+
+        # hidden state is a tuple with all the hidden layer outputs from the attention,
+        # We are only interested in the last hidden layer and the last token
+        logits = self.classifier(outputs.hidden_states[-1][:, -1, :])
+        outputs.logits = logits
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                logits, labels
+            )  # labels [batch_size], logits [batch_size x num_classes]
+            # print("LOSS", loss)
+            outputs["loss"] = loss
+
+        print("===================")
+        print(loss)
+        print("===================")
+        # print("OUTPUTS KEY" ,outputs.keys())
+        return outputs
+
+
+class MyCustomConfig(PretrainedConfig):
+    model_type = "mcq_hf_model"
+
+    def __init__(self, name_model="apple/OpenELM-450M-Instruct", **kwargs):
+        super().__init__(**kwargs)
+        self.name_model = name_model
+
+
+class MCQHFModel(PreTrainedModel):
+    config_class = MyCustomConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = MCQModel(config.name_model)
+        self.lm_head = None
+
+    def forward(self, input_ids, attention_mask=None, labels=None,
+        position_ids = None,
+        past_key_values = None,
+        inputs_embeds = None,
+        use_cache  = None,
+        output_attentions  = None,
+        output_hidden_states = None,
+        return_dict  = None,
+        cache_position = None):
+
+        return self.model(input_ids, labels=labels, attention_mask=attention_mask, position_ids= position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds,
+                            output_attentions=output_attentions, cache_position=cache_position, use_cache=use_cache, return_dict=return_dict, output_hidden_states=output_hidden_states)

model.safetensors

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