Upload modeling_gestr_jepa.py with huggingface_hub

modeling_gestr_jepa.py

@@ -0,0 +1,44 @@
+import torch
+import torch.nn as nn
+from transformers import PretrainedConfig, PreTrainedModel
+
+# --- 1. Configuration Class ---
+class GestrJEPAConfig(PretrainedConfig):
+    model_type = "gestr-jepa"
+
+    def __init__(
+        self, 
+        input_dim=16, 
+        embed_dim=64, 
+        hidden_dim=256, 
+        **kwargs
+    ):
+        self.input_dim = input_dim
+        self.embed_dim = embed_dim
+        self.hidden_dim = hidden_dim
+        super().__init__(**kwargs)
+
+# --- 2. The Model Wrapper ---
+class GestrJEPAForClassification(PreTrainedModel):
+    config_class = GestrJEPAConfig
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.encoder = nn.Sequential(
+            nn.Linear(config.input_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.hidden_dim),
+            nn.GELU(),
+            nn.Linear(config.hidden_dim, config.embed_dim),
+            nn.LayerNorm(config.embed_dim)
+        )
+        self.classifier = nn.Linear(config.embed_dim, config.num_labels)
+
+    def forward(self, sensor_values, labels=None):
+        embeddings = self.encoder(sensor_values)
+        logits = self.classifier(embeddings)
+        loss = None
+        if labels is not None:
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(logits, labels.view(-1))
+        return {"loss": loss, "logits": logits, "hidden_states": embeddings}