""" 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))}, }