checking if other model thigns are needed

config.json

@@ -141,7 +141,7 @@
   },
   "layer_norm_eps": 1e-12,
   "max_position_embeddings": 512,
-  "model_type": "stacked_bert",
+  "model_type": "lang_detect",
   "num_attention_heads": 8,
   "num_hidden_layers": 8,
   "pad_token_id": 0,

configuration_lang.py

@@ -3,7 +3,7 @@ import torch
 
 
 class ImpressoConfig(PretrainedConfig):
-    model_type = "stacked_bert"
+    model_type = "lang_detect"
 
     def __init__(
         self,

modeling_lang.py

@@ -1,10 +1,7 @@
-from transformers.modeling_outputs import TokenClassifierOutput
 import torch
 import torch.nn as nn
-from transformers import PreTrainedModel, AutoModel, AutoConfig, BertConfig
-from torch.nn import CrossEntropyLoss
-from typing import Optional, Tuple, Union
-import logging, json, os
+from transformers import PreTrainedModel
+import logging
 import floret
 from .configuration_lang import ImpressoConfig
 
@@ -26,9 +23,6 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
 
     #
     def forward(self, input_ids, attention_mask=None, **kwargs):
-        # print(
-        #     f"Check if it arrives here: {input_ids}, ---, {type(input_ids)} ----- {type(self.model_floret)}"
-        # )
         if isinstance(input_ids, str):
             # If the input is a single string, make it a list for floret
             texts = [input_ids]
@@ -37,13 +31,11 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
         else:
             raise ValueError(f"Unexpected input type: {type(input_ids)}")
 
-        # Use the SafeFloretWrapper to get predictions
         predictions, probabilities = self.model_floret.predict(texts, k=1)
-        # print(f"Predictions: {predictions}, Probabilities: {probabilities}")
         return (
             predictions,
             probabilities,
-        )  # Dummy tensor with shape (batch_size, num_classes)
+        )
 
     def state_dict(self, *args, **kwargs):
         # Return an empty state dictionary
@@ -53,21 +45,6 @@ class ExtendedMultitaskModelForTokenClassification(PreTrainedModel):
         # Ignore loading since there are no parameters
         pass
 
-        # print("Ignoring state_dict since model has no parameters.")
-
-    # def get_floret_model(self):
-    #     return self.model_floret
-
-    # def get_extended_attention_mask(
-    #     self, attention_mask, input_shape, device=None, dtype=torch.float
-    # ):
-    #     if attention_mask is None:
-    #         attention_mask = torch.ones(input_shape, device=device)
-    #     extended_attention_mask = attention_mask[:, None, None, :]
-    #     extended_attention_mask = extended_attention_mask.to(dtype=dtype)
-    #     extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
-    #     return extended_attention_mask
-
     @property
     def device(self):
         return next(self.parameters()).device