Upload modeling_te3s_head.py with huggingface_hub

modeling_te3s_head.py

@@ -1,43 +1,50 @@
-from __future__ import annotations
-
 import torch
 import torch.nn as nn
-from transformers import PreTrainedModel, PretrainedConfig
+from typing import Optional
+from transformers.modeling_utils import PreTrainedModel
+from transformers.configuration_utils import PretrainedConfig
 
 
-class TextEmbedding3SmallSentimentHead(PreTrainedModel):
-    """Lightweight sentiment head for 1536-d OpenAI embeddings.
+class TextEmbedding3SmallSentimentHeadConfig(PretrainedConfig):
+    model_type = "sentiment-head"
 
-    Expects config.json fields:
-      - input_dim (int, default 1536)
-      - hidden_dim (int, default 512; use 0 for linear-only)
-      - dropout (float, default 0.2)
-      - num_labels (int, default 3)
-    """
+    def __init__(
+        self,
+        input_dim: int = 1536,
+        hidden_dim: int = 512,
+        dropout: float = 0.2,
+        num_labels: int = 3,
+        **kwargs,
+    ) -> None:
+        super().__init__(**kwargs)
+        self.input_dim = int(input_dim)
+        self.hidden_dim = int(hidden_dim)
+        self.dropout = float(dropout)
+        self.num_labels = int(num_labels)
+
+
+class TextEmbedding3SmallSentimentHead(PreTrainedModel):
+    config_class = TextEmbedding3SmallSentimentHeadConfig
 
-    def __init__(self, config: PretrainedConfig) -> None:
+    def __init__(self, config: TextEmbedding3SmallSentimentHeadConfig) -> None:
         super().__init__(config)
-        input_dim = int(getattr(config, "input_dim", 1536))
-        hidden_dim = int(getattr(config, "hidden_dim", 512))
-        dropout = float(getattr(config, "dropout", 0.2))
-        num_labels = int(getattr(config, "num_labels", 3))
 
-        if hidden_dim and hidden_dim > 0:
+        if config.hidden_dim and config.hidden_dim > 0:
             self.net = nn.Sequential(
-                nn.Linear(input_dim, hidden_dim),
+                nn.Linear(config.input_dim, config.hidden_dim),
                 nn.ReLU(),
-                nn.Dropout(p=dropout),
-                nn.Linear(hidden_dim, num_labels),
+                nn.Dropout(p=config.dropout),
+                nn.Linear(config.hidden_dim, config.num_labels),
             )
         else:
-            self.net = nn.Linear(input_dim, num_labels)
+            self.net = nn.Linear(config.input_dim, config.num_labels)
 
         self.post_init()
 
     def forward(
         self,
         inputs_embeds: torch.FloatTensor,
-        labels: torch.LongTensor | None = None,
+        labels: Optional[torch.LongTensor] = None,
         **kwargs,
     ):
         logits = self.net(inputs_embeds)