Upload folder using huggingface_hub

modeling_havelock.py

@@ -1,75 +1,58 @@
-"""Demo: score texts on the oral–literate spectrum."""
-
 import torch
-from transformers import AutoModel, AutoTokenizer
+import torch.nn as nn
+from transformers import AutoModel, ModernBertConfig, ModernBertModel, PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import SequenceClassifierOutput
+
+
+class HavelockOralityConfig(PretrainedConfig):
+    model_type = "havelock-orality-regressor"
+
+    def __init__(self, dropout: float = 0.1, **kwargs):
+        super().__init__(**kwargs)
+        self.dropout = dropout
+
 
-from estimators.defaults import resolve
+class HavelockOralityRegressor(PreTrainedModel):
+    config_class = HavelockOralityConfig
 
-TEXTS = {
-    "Rubio (speech)": (
-        "National security, which this conference is largely about, is not merely series of technical questions – "
-        "how much we spend on defense or where, how we deploy it, these are important questions. They are. "
-        "But they are not the fundamental one. "
-        "The fundamental question we must answer at the outset is what exactly are we defending, "
-        "because armies do not fight for abstractions. "
-        "Armies fight for a people; armies fight for a nation. Armies fight for a way of life."
-    ),
-    "Academic (hedged)": (
-        "The relationship between institutional capacity and policy outcomes remains contested in the literature. "
-        "While some scholars have argued that bureaucratic autonomy is a necessary condition for effective implementation, "
-        "others have suggested that the causal mechanism operates primarily through elite consensus rather than "
-        "organizational structure per se. The evidence presented here, though preliminary, is broadly consistent "
-        "with the latter interpretation."
-    ),
-    "Reddit (casual)": (
-        "ok so I just tried this and honestly? it's not bad. like yeah the UI is kinda janky and "
-        "it crashed twice on me but once you get past that the actual features are solid. "
-        "would I recommend it? ehh maybe wait for the next update but keep an eye on it for sure"
-    ),
-    "Legal (formal)": (
-        "Notwithstanding any provision to the contrary contained herein, the obligations of the parties "
-        "under this Agreement shall survive termination to the extent necessary to effectuate the purposes "
-        "of the indemnification, confidentiality, and limitation of liability provisions set forth in "
-        "Sections 7, 9, and 12 respectively."
-    ),
-    "Homer (oral epic)": (
-        "Sing, O goddess, the anger of Achilles son of Peleus, that brought countless ills upon the Achaeans. "
-        "Many a brave soul did it send hurrying down to Hades, and many a hero did it yield a prey to dogs "
-        "and vultures, for so were the counsels of Jove fulfilled from the day on which the son of Atreus, "
-        "king of men, and great Achilles, first fell out with one another."
-    ),
-}
+    def __init__(self, config, backbone=None):
+        super().__init__(config)
+        if backbone is not None:
+            self.backbone = backbone
+        else:
+            backbone_config = ModernBertConfig.from_dict(config.to_dict())
+            self.backbone = ModernBertModel(backbone_config)
 
+        self.dropout = nn.Dropout(config.dropout)
+        self.regressor = nn.Linear(config.hidden_size, 1)
+        self.post_init()
 
-def main():
-    model_name = resolve("orality")["hub_model"]
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    @classmethod
+    def from_backbone(cls, model_name: str, dropout: float = 0.1) -> "HavelockOralityRegressor":
+        backbone = AutoModel.from_pretrained(model_name)
+        config = HavelockOralityConfig(
+            dropout=dropout,
+            **backbone.config.to_dict(),
+        )
+        return cls(config, backbone=backbone)
 
-    print(f"Loading {model_name} on {device}...")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    def _pool(self, last_hidden_state: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+        mask = attention_mask.unsqueeze(-1).float()
+        return (last_hidden_state * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
 
-    attn_impl = "flash_attention_2" if torch.cuda.is_available() else "eager"
-    model = AutoModel.from_pretrained(
-        model_name,
-        trust_remote_code=True,
-        attn_implementation=attn_impl,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    ).to(device)
-    model.eval()
+    def forward(self, input_ids, attention_mask=None, labels=None, **kwargs):
+        outputs = self.backbone(input_ids=input_ids, attention_mask=attention_mask)
 
-    print(f"\n{'Text':<25} {'Score':>6}  {'Bar'}")
-    print("=" * 65)
+        if hasattr(outputs, "pooler_output") and outputs.pooler_output is not None:
+            pooled = outputs.pooler_output
+        else:
+            pooled = self._pool(outputs.last_hidden_state, attention_mask)
 
-    for label, text in TEXTS.items():
-        inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
-        inputs = {k: v.to(device) for k, v in inputs.items()}
-        with torch.no_grad():
-            output = model(**inputs)
-            score = output.logits.squeeze().item()
-        score = max(0.0, min(1.0, score))
-        bar = "█" * int(score * 30) + "░" * (30 - int(score * 30))
-        print(f"{label:<25} {score:>5.3f}  {bar}")
+        pooled = self.dropout(pooled)
+        scores = self.regressor(pooled).squeeze(-1)
 
+        loss = None
+        if labels is not None:
+            loss = nn.MSELoss()(scores, labels)
 
-if __name__ == "__main__":
-    main()
+        return SequenceClassifierOutput(loss=loss, logits=scores)