Initial commit including model and configuration

models.py

@@ -1,28 +1,43 @@
 from transformers.modeling_outputs import TokenClassifierOutput
 import torch
 import torch.nn as nn
-from transformers import PreTrainedModel, AutoModel, AutoConfig
+from transformers import PreTrainedModel, AutoModel, AutoConfig, BertConfig
 from torch.nn import CrossEntropyLoss
 from typing import Optional, Tuple, Union
-import logging
+import logging, json, os
+
 from .configuration_stacked import ImpressoConfig
 
 logger = logging.getLogger(__name__)
 
 
+def get_info(label_map):
+    num_token_labels_dict = {task: len(labels) for task, labels in label_map.items()}
+    return num_token_labels_dict
+
+
 class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
 
     config_class = ImpressoConfig
     _keys_to_ignore_on_load_missing = [r"position_ids"]
 
-    def __init__(self, config, num_token_labels_dict):
+    def __init__(self, config):
         super().__init__(config)
-        self.num_token_labels_dict = num_token_labels_dict
+        print("Current folder path:", os.path.dirname(os.path.abspath(__file__)))
+        # Get the directory of the current script
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        # Construct the full path to label_map.json
+        label_map_path = os.path.join(current_dir, "label_map.json")
+
+        label_map = json.load(open(label_map_path, "r"))
+        self.num_token_labels_dict = get_info(label_map)
         self.config = config
 
-        # self.bert = AutoModel.from_config(config)
+        import pdb
+
+        pdb.set_trace()
         self.bert = AutoModel.from_pretrained(
-            config.name_or_path, config=config.pretrained_config
+            config.pretrained_config["_name_or_path"], config=config.pretrained_config
         )
         if "classifier_dropout" not in config.__dict__:
             classifier_dropout = 0.1
@@ -46,7 +61,7 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
         self.token_classifiers = nn.ModuleDict(
             {
                 task: nn.Linear(config.hidden_size, num_labels)
-                for task, num_labels in num_token_labels_dict.items()
+                for task, num_labels in self.num_token_labels_dict.items()
             }
         )
 

push_to_hf.py

@@ -1,9 +1,15 @@
 import os
 import shutil
 import argparse
-from transformers import AutoTokenizer, AutoConfig, AutoModelForTokenClassification
+from transformers import (
+    AutoTokenizer,
+    AutoConfig,
+    AutoModelForTokenClassification,
+    BertConfig,
+)
 from huggingface_hub import HfApi, Repository
-import json
+
+# import json
 from .configuration_stacked import ImpressoConfig
 from .models import ExtendedMultitaskModelForTokenClassification
 import subprocess
@@ -27,16 +33,13 @@ def get_info(label_map):
 
 def push_model_to_hub(checkpoint_dir, repo_name, script_path):
     checkpoint_path = get_latest_checkpoint(checkpoint_dir)
-    label_map = json.load(open(os.path.join(checkpoint_dir, "label_map.json"), "r"))
-    num_token_labels_dict = get_info(label_map)
     config = ImpressoConfig.from_pretrained(checkpoint_path)
     config.pretrained_config = AutoConfig.from_pretrained(config.name_or_path)
     config.save_pretrained("stacked_bert")
-
     config = ImpressoConfig.from_pretrained("stacked_bert")
 
     model = ExtendedMultitaskModelForTokenClassification.from_pretrained(
-        checkpoint_path, config=config, num_token_labels_dict=num_token_labels_dict
+        checkpoint_path, config=config
     )
     tokenizer = AutoTokenizer.from_pretrained(checkpoint_path)
     local_repo_path = "./repo"