Update modeling_gptj.py

modeling_gptj.py

@@ -40,6 +40,21 @@ from transformers.utils import (
 )
 from .configuration_gptj import GPTJConfig
 
+from mltools.dmx import DmxModel
+
+
+class DmxPreTrainedModel(PreTrainedModel):
+    @classmethod
+    def from_pretrained(cls, *args, **kwargs):
+        _model = super().from_pretrained(*args, **kwargs)
+        _model = DmxModel.from_torch(
+            _model, 
+            hf=True,
+            input_names=["input_ids"],  # TODO: no hard-coding!!!
+            concrete_args=None,
+        )
+        return _model
+
 
 logger = logging.get_logger(__name__)
 
@@ -53,7 +68,6 @@ GPTJ_PRETRAINED_MODEL_ARCHIVE_LIST = [
     # See all GPT-J models at https://huggingface.co/models?filter=gptj
 ]
 
-
 def create_sinusoidal_positions(num_pos: int, dim: int) -> torch.Tensor:
     inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64) / dim))
     sinusoid_inp = torch.einsum(
@@ -365,7 +379,7 @@ class GPTJBlock(nn.Module):
         return outputs  # hidden_states, present, (attentions)
 
 
-class GPTJPreTrainedModel(PreTrainedModel):
+class GPTJPreTrainedModel(DmxPreTrainedModel):
     """
     An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
     models.