dsh-router-laya 2.1.0

package/service/laya/common.py

@@ -0,0 +1,280 @@
+"""Core model architecture, token sequence construction, and confidence estimation for laya."""
+import json
+import math
+import os
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+QTYPES = {"choice": 0, "score": 1, "noul": 2}
+QTYPE_NAMES = {v: k for k, v in QTYPES.items()}
+
+
+def serialize_state(state: Union[str, dict, list]) -> str:
+    if isinstance(state, str):
+        return state
+    return json.dumps(state, ensure_ascii=False)
+
+
+def render_criterion(value) -> str:
+    """Render one criterion value as text.
+
+    Strings pass through; anything structured (dict, list, number) becomes compact JSON, so a
+    rubric reads as JSON rather than a Python repr. Without this a dict-valued criterion
+    crashed `noul` outright and leaked `{'desc': ...}` into `choice` and `score` prompts.
+    """
+    if isinstance(value, str):
+        return value
+    return json.dumps(value, ensure_ascii=False, separators=(", ", ": "), default=str)
+
+
+def render_options(q: Dict) -> List[str]:
+    """Render option texts in label-index order. Noul is always [false, true]."""
+    t, crit = q["t"], q.get("crit")
+    if t == "choice":
+        # only None/"" mean "no description"; 0 and False are legitimate criterion values
+        return [k if v is None or v == "" else "%s: %s" % (k, render_criterion(v)) for k, v in crit.items()]
+    if t == "score":
+        return ["level %d: %s" % (i, render_criterion(c)) for i, c in enumerate(crit)]
+    crit = crit or {}
+    false_crit, true_crit = crit.get("false"), crit.get("true")
+    return [
+        "false: " + (render_criterion(false_crit) if false_crit not in (None, "") else "no, the statement does not hold"),
+        "true: " + (render_criterion(true_crit) if true_crit not in (None, "") else "yes, the statement holds"),
+    ]
+
+
+def build_sequence(
+    tok,
+    state: Union[str, dict, list],
+    q: Dict,
+    max_len: int = 512,
+    head_max_len: int = 192,
+    option_order: Optional[List[int]] = None,
+    truncate_left: bool = False,
+):
+    """Format: [CLS] <type> instructions [SEP] [MASK] opt0 [MASK] opt1 ... [SEP] state [SEP]."""
+    mask_tok = tok.mask_token
+    opts = render_options(q)
+    order = option_order if option_order is not None else list(range(len(opts)))
+    ins = str(q["ins"]).replace(mask_tok, " ")
+    head_ids = tok("%s question: %s" % (q["t"], ins), add_special_tokens=False)["input_ids"]
+    opt_ids = []
+    for i in order:
+        opt_ids.append(
+            [tok.mask_token_id]
+            + tok(" " + opts[i].replace(mask_tok, " "), add_special_tokens=False)["input_ids"][:48]
+        )
+    opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    if opt_budget < 16:
+        per = max(4, (head_max_len - 16) // max(1, len(opt_ids)))
+        opt_ids = [o[:per] for o in opt_ids]
+        opt_budget = head_max_len - sum(len(o) for o in opt_ids)
+    head_ids = head_ids[: max(8, opt_budget)]
+    ids = [tok.cls_token_id] + head_ids + [tok.sep_token_id]
+    markers = []
+    for o in opt_ids:
+        markers.append(len(ids))
+        ids.extend(o)
+    ids.append(tok.sep_token_id)
+    room = max(0, max_len - len(ids) - 1)
+    st = tok(serialize_state(state).replace(mask_tok, " "), add_special_tokens=False)["input_ids"]
+    st = st[-room:] if truncate_left else st[:room]
+    ids = ids + st + [tok.sep_token_id]
+    return ids[:max_len], [m for m in markers if m < max_len]
+
+
+class DecisionModel(nn.Module):
+    """Bidirectional transformer encoder backbone + typed decision head."""
+
+    def __init__(self, encoder: nn.Module, head_layers: int = 2, n_act: int = 2, dropout: float = 0.1):
+        super().__init__()
+        self.encoder = encoder
+        d = encoder.config.hidden_size
+        nhead = max(1, d // 64)
+        layer = nn.TransformerEncoderLayer(d, nhead, 4 * d, dropout, batch_first=True, norm_first=True)
+        self.head = nn.TransformerEncoder(layer, head_layers, enable_nested_tensor=False) if head_layers > 0 else None
+        self.type_emb = nn.Embedding(3, d)
+        self.scorer = nn.Sequential(nn.LayerNorm(d), nn.Linear(d, d), nn.GELU(), nn.Linear(d, 1))
+        self.act_head = nn.Sequential(nn.Linear(d + 4, 256), nn.GELU(), nn.Linear(256, n_act))
+        self.register_buffer("temperature", torch.ones(3))
+        self.head_checkpointing = False
+
+    def forward(self, input_ids, attention_mask, marker_pos, marker_mask, qtype, detach_encoder: bool = False):
+        h = self.encoder(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state
+        if detach_encoder:
+            h = h.detach()
+        h = h + self.type_emb(qtype)[:, None, :]
+        if self.head is not None:
+            pad = ~attention_mask.bool()
+            for layer in self.head.layers:
+                h = layer(h, src_key_padding_mask=pad)
+        idx = marker_pos.clamp(min=0)[:, :, None].expand(-1, -1, h.size(-1))
+        m = torch.gather(h, 1, idx)
+        logits = self.scorer(m).squeeze(-1).float()
+        logits = logits.masked_fill(~marker_mask, -1e4)
+
+        p = torch.softmax(logits.detach(), -1)
+        k = marker_mask.sum(-1).clamp(min=2).float()
+        ent = -(p * torch.log(p.clamp_min(1e-9))).sum(-1) / torch.log(k)
+        if p.size(-1) >= 2:
+            top2 = p.topk(2, -1).values
+        else:
+            # A single-option question has exactly one marker, so p.topk(2, ...)
+            # has nothing to select for the second slot and raises. The answer
+            # is still well-defined: softmax over one logit is 1.0 regardless of
+            # its value, so pad the missing second entry with 0.0 - that gives
+            # the act head top1 - top2 == 1.0, the same "fully decided" signal
+            # it would see for any other unambiguous top-1-vs-rest gap.
+            top1 = p.topk(1, -1).values
+            top2 = torch.cat([top1, torch.zeros_like(top1)], dim=-1)
+        feats = torch.stack([top2[:, 0], top2[:, 0] - top2[:, 1], ent, k / 255.0], -1)
+        pooled = h[:, 0].float()
+        act_logits = self.act_head(torch.cat([pooled, feats], -1))
+        return logits, act_logits
+
+
+def build_model(cfg: Dict, encoder_dir: Optional[str] = None, pretrained: bool = True) -> DecisionModel:
+    from transformers import AutoConfig, AutoModel
+
+    if not pretrained or (encoder_dir and os.path.exists(encoder_dir)):
+        ecfg = AutoConfig.from_pretrained(encoder_dir or cfg["encoder"])
+        enc = AutoModel.from_config(ecfg, attn_implementation="sdpa")
+    else:
+        enc = AutoModel.from_pretrained(cfg["encoder"], attn_implementation="sdpa")
+    return DecisionModel(enc, cfg.get("head_layers", 2), len(cfg.get("act_costs", {})) + 1)
+
+
+def proper_reward(
+    q: torch.Tensor,
+    target: torch.Tensor,
+    qtype: torch.Tensor,
+    mask: torch.Tensor,
+    w_sph: float = 0.5,
+    w_rps: float = 1.0,
+    log_floor: float = -9.21,
+) -> torch.Tensor:
+    """Strictly proper scoring rule reward: log score + spherical score + ranked probability score.
+
+    q: [..., N, K] reported distributions
+    target: [N, K] (one-hot or soft target distributions)
+    """
+    q = q * mask
+    logq = torch.log(q.clamp_min(1e-12)).clamp_min(log_floor)
+    log_score = (target * logq).sum(-1)
+    sph = (target * q).sum(-1) / q.norm(dim=-1).clamp_min(1e-9)
+    r = log_score + w_sph * sph
+    is_score = (qtype == QTYPES["score"]).float()
+    if is_score.any():
+        k = mask.sum(-1).clamp(min=2).float()
+        cdf_q = torch.cumsum(q, -1)
+        cdf_t = torch.cumsum(target, -1)
+        rps = (((cdf_q - cdf_t) ** 2) * mask).sum(-1) / (k - 1)
+        r = r - w_rps * rps * is_score
+    return r
+
+
+def td_lambda_targets(p_true: torch.Tensor, batch: Dict, lam: float = 1.0) -> torch.Tensor:
+    """TD(lambda) targets for multi-turn conversation trajectories."""
+    target = batch["target"].clone()
+    groups = batch.get("ep_group")
+    if groups is None:
+        return target
+    for g in torch.unique(groups[groups >= 0]).tolist():
+        idx = (groups == g).nonzero(as_tuple=True)[0]
+        idx = idx[torch.argsort(batch["ep_step"][idx])]
+        y = batch["target"][idx[-1], 1]
+        G = y
+        for j in range(len(idx) - 1, -1, -1):
+            if j < len(idx) - 1:
+                G = (1 - lam) * p_true[idx[j + 1]] + lam * G
+            target[idx[j], 0], target[idx[j], 1] = 1 - G, G
+    return target
+
+
+def ece_score(conf: np.ndarray, correct: np.ndarray, bins: int = 15) -> float:
+    """Expected Calibration Error across confidence bins."""
+    if len(conf) == 0:
+        return float("nan")
+    edges = np.linspace(0, 1, bins + 1)
+    e = 0.0
+    for lo, hi in zip(edges[:-1], edges[1:]):
+        sel = (conf > lo) & (conf <= hi)
+        if sel.any():
+            e += sel.mean() * abs(conf[sel].mean() - correct[sel].mean())
+    return float(e)
+
+
+def confidence_from_probs(p: np.ndarray, k: int) -> float:
+    """Normalized Shannon entropy confidence: 1 - H(p) / log(k)."""
+    if k < 2:
+        return 1.0
+    p = p[:k]
+    ent = -(p * np.log(np.clip(p, 1e-12, 1.0))).sum()
+    return float(np.clip(1.0 - ent / math.log(k), 0.0, 1.0))
+
+
+def temp_bucket(qtype: int, k: int) -> str:
+    size = "2" if k <= 2 else "3-5" if k <= 5 else "6-10" if k <= 10 else "11+"
+    return "%s:%s" % (QTYPE_NAMES[int(qtype)], size)
+
+
+# A fitted temperature below 1 sharpens the logits instead of softening them. The shipped
+# `choice:11+` bucket is 0.1006, which multiplies them ~10x: a 0.24 top probability is published as
+# 0.99, so a caller gating on confidence is told a coin flip is a certainty. No honest calibration
+# needs to sharpen this hard, so refuse to apply one that does.
+TEMP_MIN = 0.5
+TEMP_MAX = 5.0
+
+
+def clamp_temperature(t, lo: float = TEMP_MIN, hi: float = TEMP_MAX) -> float:
+    """A usable temperature: `t` confined to [lo, hi], falling back to 1.0 if it is not a number."""
+    try:
+        t = float(t)
+    except (TypeError, ValueError):
+        return 1.0
+    if t != t or t in (float("inf"), float("-inf")):    # NaN / inf
+        return 1.0
+    return min(hi, max(lo, t))
+
+
+def amp_dtype(name: Optional[str]) -> torch.dtype:
+    return torch.bfloat16 if name == "bf16" else torch.float16
+
+
+def collate_items(batch, pad_id: int):
+    items = [it for group in batch for it in group]
+    if not items:
+        return None
+    n, L = len(items), max(len(it["ids"]) for it in items)
+    kmax = max(len(it["markers"]) for it in items)
+    ids = torch.full((n, L), pad_id, dtype=torch.long)
+    att = torch.zeros((n, L), dtype=torch.long)
+    mpos = torch.zeros((n, kmax), dtype=torch.long)
+    mmask = torch.zeros((n, kmax), dtype=torch.bool)
+    has_target = any("target" in it for it in items)
+    target = torch.zeros((n, kmax), dtype=torch.float32) if has_target else None
+
+    for i, it in enumerate(items):
+        ids[i, : len(it["ids"])] = torch.tensor(it["ids"])
+        att[i, : len(it["ids"])] = 1
+        k = len(it["markers"])
+        mpos[i, :k] = torch.tensor(it["markers"])
+        mmask[i, :k] = True
+        if has_target and "target" in it:
+            target[i, : len(it["target"])] = torch.tensor(it["target"], dtype=torch.float32)
+
+    res = {
+        "input_ids": ids,
+        "attention_mask": att,
+        "marker_pos": mpos,
+        "marker_mask": mmask,
+        "qtype": torch.tensor([it["qtype"] for it in items]),
+        "label": torch.tensor([it.get("label", -1) for it in items]),
+        "meta": [{k: it[k] for k in it if k not in ("ids", "markers", "target")} for it in items],
+    }
+    if target is not None:
+        res["target"] = target
+    return res

package/service/laya/email.py

@@ -0,0 +1,90 @@
+"""Email utilities for cleaning and structuring email inputs in laya."""
+import re
+from typing import Dict, Optional
+
+_QUOTE_HEADERS = [
+    re.compile(r"^\s*On .{0,300}wrote:\s*$", re.I),
+    re.compile(r"^\s*-{2,}\s*(Original|Forwarded) Message\s*-{2,}", re.I),
+    re.compile(r"^\s*_{8,}\s*$"),
+    re.compile(r"^\s*From:\s.+$", re.I),
+]
+_SIGNATURE_MARKERS = [
+    re.compile(r"^\s*--\s*$"),
+    re.compile(r"^\s*(best|kind|warm|many thanks|thanks|thank you|regards|cheers|sincerely)[\w ,!.]*$", re.I),
+    re.compile(r"^\s*sent from my (iphone|android|mobile|ipad)", re.I),
+]
+_DISCLAIMER = re.compile(
+    r"(confidential|intended (solely )?for the (use of the )?(named )?(addressee|recipient)|"
+    r"if you (have )?received this (e-?mail|message) in error)",
+    re.I,
+)
+
+
+def clean_email_body(body: str, max_chars: int = 3000) -> str:
+    """Remove quoted email history, signatures and disclaimers to keep input focused."""
+    text = (body or "").replace("\r\n", "\n").replace("\r", "\n").replace("\\n", "\n")
+    lines = []
+    for line in text.split("\n"):
+        if any(p.match(line) for p in _QUOTE_HEADERS) and lines:
+            break
+        if line.lstrip().startswith(">"):
+            continue
+        lines.append(line.rstrip())
+    cut = len(lines)
+    for i in range(max(1, min(int(len(lines) * 0.6), len(lines) - 8)), len(lines)):
+        if len(lines[i].strip()) <= 40 and any(p.match(lines[i]) for p in _SIGNATURE_MARKERS):
+            cut = i
+            break
+    lines = lines[:cut]
+    paragraphs = [p for p in re.split(r"\n\s*\n", "\n".join(lines)) if not _DISCLAIMER.search(p)]
+    text = re.sub(r"[ \t]+", " ", "\n\n".join(p.strip() for p in paragraphs if p.strip()))
+    return text[:max_chars]
+
+
+def email_state(subject: str, body: str, sender: Optional[str] = None, clean: bool = True, **extra) -> Dict:
+    """Construct a clean state dictionary for email classification."""
+    state = {
+        "subject": (subject or "").strip(),
+        "body": clean_email_body(body) if clean else (body or ""),
+    }
+    if sender:
+        state["from"] = sender
+    state.update({k: v for k, v in extra.items() if v is not None})
+    return state
+
+
+def email_questions(categories: Optional[Dict[str, str]] = None) -> Dict:
+    """Standard pre-built questions for email triage."""
+    categories = categories or {
+        "billing": "invoices, payments, refunds",
+        "technical": "bugs, outages, integrations",
+        "sales": "pricing, demos, new purchases",
+        "security": "phishing, scams, account compromise",
+        "hr": "hiring, leave, payroll",
+        "other": "none of the above",
+    }
+    return {
+        "category": {
+            "type": "choice",
+            "instructions": "Which team should handle the email in `body`?",
+            "criteria": categories,
+        },
+        "is_spam": {
+            "type": "noul",
+            "instructions": "Is this email unsolicited spam or bulk marketing?",
+        },
+        "is_phishing": {
+            "type": "noul",
+            "instructions": "Is this email a phishing or scam attempt to steal money, credentials, or personal data?",
+            "criteria": {"true": "phishing, scam, or fraud", "false": "a legitimate email"},
+        },
+        "urgency": {
+            "type": "score",
+            "instructions": "How urgent is the request in `body`?",
+            "criteria": ["no time pressure", "needs attention soon", "blocking issue or hard deadline"],
+        },
+        "needs_reply": {
+            "type": "noul",
+            "instructions": "Does the sender expect a reply?",
+        },
+    }