Blame for modeling_scribelm.py

"""
Custom Hugging Face PreTrainedConfig and PreTrainedModel for `scribelm`.

Architecture: `ScribeLMForAuthoring` (~134.8M parameters)
- Backbone: Intact 30-layer Llama/SmolLM2 Causal Transformer (`hidden_size=576`, `num_attention_heads=9`,
  `num_key_value_heads=3`, `intermediate_size=1536`, `vocab_size=49152`) initialized from
  `HuggingFaceTB/SmolLM2-135M-Instruct` with bottom-20-layer grammar preservation and top-10-layer
  Hemanth HM style adaptation.
- Joint Multi-Task Heads:
  1. `lm_head`: Causal language modeling head generating Hemanth HM blog posts, rewrites, and social threads.
  2. `voice_head`: Calibrated `p_hemanth` logit scoring authentic practitioner voice vs corporate/LLM fluff.
  3. `style_head`: 5-axis style telemetry head predicting `(directness, hook_strength, zero_fluff, code_first, cadence)`.
  4. `task_head`: 4-way writing task router `(blog, rewrite, social, critic)`.
"""

from __future__ import annotations

from typing import Any, Dict, List, Optional

import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig, PreTrainedModel

_TF_V5 = int(transformers.__version__.split(".")[0]) >= 5

TASK_NAMES = ["blog", "rewrite", "social", "critic"]
TASK_TO_ID = {name: i for i, name in enumerate(TASK_NAMES)}
STYLE_AXES = ["directness", "hook_strength", "zero_fluff", "code_first", "cadence"]

HEMANTH_SYSTEM_PROMPT = (
    "You are Hemanth HM (h3manth.com/scribe). Write direct, punchy, zero-fluff practitioner prose "
    "with short paragraphs, concrete metrics, and bullet points."
)


class ScribeLMConfig(PretrainedConfig):
    model_type = "scribelm"

    def __init__(
        self,
        base_model_name: str = "HuggingFaceTB/SmolLM2-135M-Instruct",
        vocab_size: int = 49152,
        hidden_size: int = 576,
        intermediate_size: int = 1536,
        num_hidden_layers: int = 30,
        num_attention_heads: int = 9,
        num_key_value_heads: int = 3,
        hidden_act: str = "silu",
        max_position_embeddings: int = 2048,
        rms_norm_eps: float = 1e-5,
        rope_theta: float = 100000.0,
        num_tasks: int = 4,
        num_style_axes: int = 5,
        dropout: float = 0.1,
        freeze_bottom_layers: int = 20,
        **kwargs: Any,
    ) -> None:
        super().__init__(**kwargs)
        self.base_model_name = base_model_name
        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.intermediate_size = intermediate_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_key_value_heads = num_key_value_heads
        self.hidden_act = hidden_act
        self.max_position_embeddings = max_position_embeddings
        self.rms_norm_eps = rms_norm_eps
        self.rope_theta = rope_theta
        self.num_tasks = num_tasks
        self.num_style_axes = num_style_axes
        self.dropout = dropout
        self.freeze_bottom_layers = freeze_bottom_layers


class ScribeLMForAuthoring(PreTrainedModel):
    config_class = ScribeLMConfig
    _tied_weights_keys = (
        {"causal_lm.lm_head.weight": "causal_lm.model.embed_tokens.weight"}
        if _TF_V5
        else ["causal_lm.lm_head.weight"]
    )

    def __init__(self, config: ScribeLMConfig, preload_pretrained: bool = False) -> None:
        super().__init__(config)
        self.config = config
        if _TF_V5:
            self.all_tied_weights_keys = {"causal_lm.lm_head.weight": "causal_lm.model.embed_tokens.weight"}

        llama_cfg = AutoConfig.from_pretrained(config.base_model_name)
        llama_cfg.num_hidden_layers = config.num_hidden_layers
        llama_cfg._attn_implementation = "eager"

        if preload_pretrained:
            full_lm = AutoModelForCausalLM.from_pretrained(config.base_model_name)
            full_lm.config._attn_implementation = "eager"
            self.causal_lm = full_lm.float()
        else:
            self.causal_lm = AutoModelForCausalLM.from_config(llama_cfg).float()

        self.causal_lm = self.causal_lm.float()
        self.causal_lm.lm_head.weight = self.causal_lm.model.embed_tokens.weight

        h = config.hidden_size
        self.dropout = nn.Dropout(config.dropout)
        self.voice_head = nn.Sequential(
            nn.Linear(h, h // 2),
            nn.SiLU(),
            nn.Linear(h // 2, 1),
        )
        self.style_head = nn.Sequential(
            nn.Linear(h, h // 2),
            nn.SiLU(),
            nn.Linear(h // 2, config.num_style_axes),
        )
        self.task_head = nn.Linear(h, config.num_tasks)

    def tie_weights(self, **kwargs: Any) -> None:
        self.causal_lm.lm_head.weight = self.causal_lm.model.embed_tokens.weight

    def freeze_lower_layers(self, num_freeze: Optional[int] = None) -> None:
        n_f = self.config.freeze_bottom_layers if num_freeze is None else num_freeze
        self.causal_lm.model.embed_tokens.weight.requires_grad = False
        for name, p in self.causal_lm.named_parameters():
            if "model.layers." in name:
                layer_idx = int(name.split("model.layers.")[1].split(".")[0])
                if layer_idx < n_f:
                    p.requires_grad = False

    def forward(
        self,
        input_ids: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        voice_labels: Optional[torch.Tensor] = None,
        style_labels: Optional[torch.Tensor] = None,
        task_labels: Optional[torch.Tensor] = None,
    ) -> Dict[str, torch.Tensor]:
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)

        out = self.causal_lm(
            input_ids=input_ids,
            attention_mask=attention_mask,
            use_cache=False,
            output_hidden_states=True,
        )
        logits = out.logits
        last_hidden = out.hidden_states[-1]

        mask_f = attention_mask.unsqueeze(-1).to(last_hidden.dtype)
        pooled = (last_hidden * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp(min=1.0)
        pooled_drop = self.dropout(pooled)

        voice_logit = self.voice_head(pooled_drop).squeeze(-1)
        style_pred = torch.sigmoid(self.style_head(pooled_drop))
        task_logits = self.task_head(pooled_drop)

        result: Dict[str, torch.Tensor] = {
            "logits": logits,
            "next_token_logits": logits[:, -1, :],
            "voice_logit": voice_logit,
            "p_hemanth": torch.sigmoid(voice_logit),
            "style_scores": style_pred,
            "task_logits": task_logits,
        }

        total_loss = None
        if labels is not None:
            shift_logits = logits[:, :-1, :].contiguous()
            shift_labels = labels[:, 1:].contiguous()
            lm_loss = F.cross_entropy(
                shift_logits.view(-1, shift_logits.size(-1)),
                shift_labels.view(-1),
                ignore_index=-100,
            )
            result["lm_loss"] = lm_loss
            total_loss = lm_loss

        if voice_labels is not None:
            v_loss = F.binary_cross_entropy_with_logits(
                voice_logit.view(-1),
                voice_labels.view(-1).to(voice_logit.dtype),
            )
            result["voice_loss"] = v_loss
            total_loss = v_loss if total_loss is None else total_loss + 0.35 * v_loss

        if style_labels is not None:
            s_loss = F.smooth_l1_loss(style_pred, style_labels.to(style_pred.dtype))
            result["style_loss"] = s_loss
            total_loss = s_loss if total_loss is None else total_loss + 0.45 * s_loss

        if task_labels is not None:
            t_loss = F.cross_entropy(task_logits, task_labels.view(-1).long())
            result["task_loss"] = t_loss
            total_loss = t_loss if total_loss is None else total_loss + 0.20 * t_loss

        if total_loss is not None:
            result["loss"] = total_loss

        return result

    @torch.no_grad()
    def generate_text(
        self,
        tokenizer: Any,
        user_prompt: str,
        max_new_tokens: int = 120,
        repetition_penalty: float = 1.12,
    ) -> Dict[str, Any]:
        self.eval()
        device = next(self.parameters()).device
        messages = [
            {"role": "system", "content": HEMANTH_SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ]
        chat_prompt = tokenizer.apply_chat_template(
            messages,
            tokenize=False,
            add_generation_prompt=True,
        )
        enc = tokenizer(
            chat_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=256,
        ).to(device)

        gen_ids = self.causal_lm.generate(
            **enc,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            repetition_penalty=repetition_penalty,
        )
        new_tokens = gen_ids[0][enc["input_ids"].shape[1] :]
        text_out = tokenizer.decode(new_tokens, skip_special_tokens=True).strip()

        full_mask = torch.ones_like(gen_ids)
        scored_out = self.forward(
            input_ids=gen_ids,
            attention_mask=full_mask,
        )
        p_hemanth = float(scored_out["p_hemanth"][0].item())
        style_vec = [round(float(x), 4) for x in scored_out["style_scores"][0].tolist()]
        task_probs = torch.softmax(scored_out["task_logits"][0], dim=-1).tolist()
        pred_task_idx = int(torch.argmax(scored_out["task_logits"][0]).item())

        return {
            "text": text_out,
            "p_hemanth": round(p_hemanth, 4),
            "predicted_task": TASK_NAMES[pred_task_idx],
            "task_probs": {TASK_NAMES[i]: round(float(task_probs[i]), 4) for i in range(len(TASK_NAMES))},
            "style_scores": {STYLE_AXES[i]: style_vec[i] for i in range(len(STYLE_AXES))},
        }