| import torch |
| from transformers import PreTrainedModel |
| from transformers.modeling_outputs import ModelOutput |
| from dataclasses import dataclass, field |
| from typing import Optional |
|
|
| from .configuration_domsense import SystemOneConfig |
|
|
| |
| from .src_models.system_one_model import SystemOneModel as _SystemOneCore |
|
|
|
|
| @dataclass |
| class DomSenseOutput(ModelOutput): |
| """DomSense 决策模型的前向输出。 |
| |
| 同时兼顾「选择」与「MDP 前瞻」产物，便于上层 / 用户直接消费。 |
| """ |
|
|
| loss: Optional[torch.FloatTensor] = None |
| loss_calib: torch.FloatTensor | None = None |
| loss_q: torch.FloatTensor | None = None |
| loss_decision: torch.FloatTensor | None = None |
| choices: torch.LongTensor | None = None |
| choice_probs: torch.FloatTensor | None = None |
| transition_probs: torch.FloatTensor | None = None |
| expected_returns: torch.FloatTensor | None = None |
| q_values: torch.FloatTensor | None = None |
| confidences: torch.FloatTensor | None = None |
| q_history: Optional[list] = None |
| rewards: torch.FloatTensor | None = None |
| embeddings: torch.FloatTensor | None = None |
| action_embeddings: torch.FloatTensor | None = None |
|
|
|
|
| class SystemOneModelForDecision(PreTrainedModel): |
| """Hugging Face 标准接口包装：把 DomSense 决策内核暴露成 PreTrainedModel。 |
| |
| - ``from_pretrained`` 可加载：仓库内 ``config.json`` + ``model.safetensors`` 或 |
| 原始 ``model_state_dict``。 |
| - ``save_pretrained`` 可导出：结构与权重统一为 HF 标准格式。 |
| - 内核权重名与原始训练 checkpoint 完全一致（顶层前缀 ``core.``），因此 |
| 转换/反向转换都无损。 |
| """ |
|
|
| config_class = SystemOneConfig |
| base_model_prefix = "core" |
| main_input_name = "input_texts" |
|
|
| def __init__(self, config: SystemOneConfig): |
| super().__init__(config) |
| self.config = config |
|
|
| |
| core_cfg = config.to_dict() |
| self.core = _SystemOneCore(core_cfg) |
|
|
| |
| |
| self.post_init() |
|
|
| |
| |
| |
| def set_encoder_backend(self, backend: str): |
| """切换编码器后端：'lightweight'（内置字符嵌入）或 'transformers'（distilbert）。 |
| |
| 注意：切换后端意味着编码器结构与权重不同，需用对应 checkpoint 加载； |
| 本仓库附带的权重为 lightweight 后端。切换仅在你有对应 distilbert 权重时使用。 |
| """ |
| raise NotImplementedError( |
| "编码器后端切换需要对应权重；本仓库自带 lightweight 后端权重，" |
| "若需 transformers(distilbert) 后端请自行加载对应 checkpoint。" |
| ) |
|
|
| |
| |
| |
| def forward( |
| self, |
| input_texts=None, |
| action_texts=None, |
| return_mdp_details: bool = False, |
| num_valid_actions=None, |
| labels=None, |
| **kwargs, |
| ): |
| """前向传播。 |
| |
| Args: |
| input_texts: List[str] 场景文本。 |
| action_texts: list of list[str]，每个样本的选项文本。 |
| return_mdp_details: 是否返回 MDP 内部细节。 |
| num_valid_actions: (batch,) 有效动作数。 |
| labels: 可选。结构化 label dict，用于计算 joint loss（尚未接入）。 |
| """ |
| if input_texts is None and "texts" in kwargs: |
| input_texts = kwargs.pop("texts") |
| if input_texts is None: |
| raise ValueError("input_texts 不能为空") |
|
|
| out = self.core.forward( |
| input_texts, |
| action_texts=action_texts, |
| num_valid_actions=num_valid_actions, |
| return_mdp_details=return_mdp_details, |
| ) |
|
|
| action_emb = None |
| if action_texts is not None: |
| action_emb, _ = self.core.encode_options(action_texts) |
|
|
| return DomSenseOutput( |
| choices=out["choices"], |
| choice_probs=out["choice_probs"], |
| transition_probs=out["transition_probs"], |
| expected_returns=out["expected_returns"], |
| q_values=out["q_values"], |
| confidences=out["confidences"], |
| q_history=out.get("q_history", []), |
| rewards=out.get("rewards"), |
| embeddings=out["embeddings"], |
| action_embeddings=action_emb, |
| ) |
|
|
| def generate_compact_output(self, texts, action_texts=None, num_valid_actions=None): |
| """生成紧凑决策输出（与训练/服务层输出契约一致）。""" |
| return self.core.generate_compact_output( |
| texts, action_texts=action_texts, num_valid_actions=num_valid_actions |
| ) |
|
|
| def count_parameters(self): |
| return self.core.count_parameters() |
|
|
| |
| |
| |
| def _init_weights(self, module): |
| pass |
|
|
