Upload modeling_te3s_head.py with huggingface_hub

modeling_te3s_head.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from transformers import PreTrainedModel, PretrainedConfig
+
+
+class TextEmbedding3SmallSentimentHead(PreTrainedModel):
+    """Lightweight sentiment head for 1536-d OpenAI embeddings.
+
+    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, config: PretrainedConfig) -> 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:
+            self.net = nn.Sequential(
+                nn.Linear(input_dim, hidden_dim),
+                nn.ReLU(),
+                nn.Dropout(p=dropout),
+                nn.Linear(hidden_dim, num_labels),
+            )
+        else:
+            self.net = nn.Linear(input_dim, num_labels)
+
+        self.post_init()
+
+    def forward(
+        self,
+        inputs_embeds: torch.FloatTensor,
+        labels: torch.LongTensor | None = None,
+        **kwargs,
+    ):
+        logits = self.net(inputs_embeds)
+        loss = None
+        if labels is not None:
+            loss = nn.CrossEntropyLoss()(logits, labels)
+        return {"logits": logits, "loss": loss}
+
+