Release final shared-prompt-gating checkpoint

Replace the previous checkpoint with the selected final model and a common model-card format across the six-model family.

README.md

@@ -1,130 +1,121 @@
 ---
 library_name: transformers
-base_model: Skywork/Skywork-Reward-V2-Llama-3.1-8B
 pipeline_tag: text-classification
-language:
-- en
-datasets:
-- RLHFlow/UltraFeedback-preference-standard
-- allenai/reward-bench
+base_model: Skywork/Skywork-Reward-V2-Llama-3.1-8B
 tags:
 - reward-model
-- preference-modeling
-- rlhf
-- multi-domain
-- coherence
-- commonsense
-- empathy
-- multicultural
+- preference-model
+- multidomain
+- shared-prompt-gating
 ---
 
-# Multi-Domain Reward Model Skywork Llama-3.1-8B-Instruct
-
-This is a multi-domain reward model built from
-[`Skywork/Skywork-Reward-V2-Llama-3.1-8B`](https://huggingface.co/Skywork/Skywork-Reward-V2-Llama-3.1-8B).
-It combines 23 fine-grained regression objectives across coherence, commonsense,
-empathy, and multicultural response quality with a prompt-conditioned gating
-network that produces a single preference score.
-
-The checkpoint was packaged with the custom `RewardModelWithGating`
-architecture used in the Multi-Domain Reward Model project.
-
-Project repository: [`Mario-RC/multi-domain-reward-model`](https://github.com/Mario-RC/multi-domain-reward-model).
-
-## Intended Use
-
-Use this model to score and compare assistant responses when the evaluation
-should account for multiple quality dimensions rather than a single generic
-helpfulness score. The primary use case is reward modeling or offline response
-ranking for chat-style data.
-
-## Training Data
+# Multi-Domain Reward Model — Skywork Llama 3.1 8B
 
-The model uses multi-objective scoring and preference data from:
+This is the final Multi-Domain Reward Model checkpoint built on
+[`Skywork/Skywork-Reward-V2-Llama-3.1-8B`](https://huggingface.co/Skywork/Skywork-Reward-V2-Llama-3.1-8B). It scores responses across
+coherence, commonsense, empathy, and multicultural understanding using 23 fine-grained attributes
+and one prompt-conditioned gate shared by both responses in a preference pair.
 
-- [`multidomain_data_scoring`](https://github.com/mestecha/multidomain_data_scoring)
-- [`RLHFlow/UltraFeedback-preference-standard`](https://huggingface.co/datasets/RLHFlow/UltraFeedback-preference-standard)
-- [`allenai/reward-bench`](https://huggingface.co/datasets/allenai/reward-bench) for evaluation
+## Released model family
 
-## Evaluation
+All checkpoints use the same held-out set of 15,489 preference pairs. Confidence intervals are 95%
+prompt-cluster bootstrap intervals.
 
-Preference accuracy by domain:
+| Model | Base reward model | Test accuracy | 95% CI |
+|---|---|---:|---:|
+| [FsfairX Gemma 2 9B](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-gemma-2-9b-it) | [sfairXC/FsfairX-Gemma2-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-Gemma2-RM-v0.1) | **88.01%** | [87.43, 88.59] |
+| [Skywork Qwen 3 8B](https://huggingface.co/mario-rc/multi-domain-rm-skywork-qwen-3-8b-it) | [Skywork/Skywork-Reward-V2-Qwen3-8B](https://huggingface.co/Skywork/Skywork-Reward-V2-Qwen3-8B) | **87.82%** | [87.24, 88.40] |
+| [FsfairX Llama 3 8B](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-llama-3-8b-it) | [sfairXC/FsfairX-LLaMA3-RM-v0.1](https://huggingface.co/sfairXC/FsfairX-LLaMA3-RM-v0.1) | **86.86%** | [86.24, 87.48] |
+| [Skywork Llama 3.1 8B](https://huggingface.co/mario-rc/multi-domain-rm-skywork-llama-3.1-8b-it) | [Skywork/Skywork-Reward-V2-Llama-3.1-8B](https://huggingface.co/Skywork/Skywork-Reward-V2-Llama-3.1-8B) | **86.82%** | [86.20, 87.43] |
+| [Mistral 7B](https://huggingface.co/mario-rc/multi-domain-rm-mistral-7b-it) | [weqweasdas/RM-Mistral-7B](https://huggingface.co/weqweasdas/RM-Mistral-7B) | **84.41%** | [83.75, 85.06] |
+| [Qwen 3 Nemotron 8B](https://huggingface.co/mario-rc/multi-domain-rm-qwen-3-nemotron-8b-it) | [nvidia/Qwen3-Nemotron-8B-BRRM](https://huggingface.co/nvidia/Qwen3-Nemotron-8B-BRRM) | **83.65%** | [82.96, 84.30] |
 
-| Domain | Accuracy (%) |
-| :--- | ---: |
-| Coherence | 87.5270 |
-| Commonsense | 98.0365 |
-| Empathy | 96.2669 |
-| Multicultural | 85.2487 |
+## This checkpoint
 
-## Hugging Face Models
+| Metric | Result |
+|---|---:|
+| Held-out preference accuracy | **86.82%** |
+| 95% confidence interval | [86.20, 87.43] |
+| Coherence accuracy | 75.59% |
+| Commonsense accuracy | 97.39% |
+| Empathy accuracy | 92.72% |
+| Multicultural accuracy | 74.99% |
 
-The packaged multi-domain reward models are available on Hugging Face under the `mario-rc` namespace:
+## Usage
 
-| Model | Base reward model | Preference accuracy (%) | Scoring Spearman |
-| --- | --- | :---: | :---: |
-| [`multi-domain-rm-fsfairx-gemma-2-9b-it`](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-gemma-2-9b-it) | [`sfairXC/FsfairX-Gemma2-RM-v0.1`](https://huggingface.co/sfairXC/FsfairX-Gemma2-RM-v0.1) | 92.56 | 0.7346 |
-| [`multi-domain-rm-fsfairx-llama-3-8b-it`](https://huggingface.co/mario-rc/multi-domain-rm-fsfairx-llama-3-8b-it) | [`sfairXC/FsfairX-LLaMA3-RM-v0.1`](https://huggingface.co/sfairXC/FsfairX-LLaMA3-RM-v0.1) | 89.87 | 0.7108 |
-| [`multi-domain-rm-mistral-7b-it`](https://huggingface.co/mario-rc/multi-domain-rm-mistral-7b-it) | [`weqweasdas/RM-Mistral-7B`](https://huggingface.co/weqweasdas/RM-Mistral-7B) | 91.69 | 0.6710 |
-| [`multi-domain-rm-qwen-3-nemotron-8b-it`](https://huggingface.co/mario-rc/multi-domain-rm-qwen-3-nemotron-8b-it) | [`nvidia/Qwen3-Nemotron-8B-BRRM`](https://huggingface.co/nvidia/Qwen3-Nemotron-8B-BRRM) | 93.36 | 0.6704 |
-| [`multi-domain-rm-skywork-llama-3.1-8b-it`](https://huggingface.co/mario-rc/multi-domain-rm-skywork-llama-3.1-8b-it) | [`Skywork/Skywork-Reward-V2-Llama-3.1-8B`](https://huggingface.co/Skywork/Skywork-Reward-V2-Llama-3.1-8B) | 92.69 | 0.7264 |
-| [`multi-domain-rm-skywork-qwen-3-8b-it`](https://huggingface.co/mario-rc/multi-domain-rm-skywork-qwen-3-8b-it) | [`Skywork/Skywork-Reward-V2-Qwen3-8B`](https://huggingface.co/Skywork/Skywork-Reward-V2-Qwen3-8B) | 95.72 | 0.7156 |
-
-## Usage Example
-
-This checkpoint uses the project's custom `RewardModelWithGating` class. Run the
-example from an environment where `multidomain_model/modeling_custom.py` is
-importable.
+The repository includes custom Transformers code, so `trust_remote_code=True` is required. Compute
+the gate once from the prompt and reuse that tensor when scoring both complete candidates.
 
 ```python
 import torch
-from transformers import AutoTokenizer
-from modeling_custom import RewardModelWithGating
-
-model_id = "mario-rc/multi-domain-rm-skywork-llama-3.1-8b-it"
-dtype = torch.bfloat16 if torch.cuda.is_available() else torch.float32
-device_map = {"": 0} if torch.cuda.is_available() else None
-
-tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)
-model = RewardModelWithGating.from_pretrained(
-    model_id,
-    device_map=device_map,
-    dtype=dtype,
+from transformers import AutoModel, AutoTokenizer
+
+repo_id = "mario-rc/multi-domain-rm-skywork-llama-3.1-8b-it"
+tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+model = AutoModel.from_pretrained(
+    repo_id,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    trust_remote_code=True,
 ).eval()
-device = next(model.parameters()).device
 
-messages = [
-    {"role": "user", "content": "I failed an important exam and feel awful."},
-    {"role": "assistant", "content": "I'm sorry. That is a hard setback, but it does not define your ability. Take a little time to recover, then we can make a concrete study plan for the next attempt."},
-]
-
-encoded = tokenizer.apply_chat_template(
-    messages,
+prompt = [{"role": "user", "content": "How can I support a friend who feels excluded?"}]
+chosen = prompt + [{
+    "role": "assistant",
+    "content": "Listen without judging, validate how they feel, and ask what support would help.",
+}]
+rejected = prompt + [{"role": "assistant", "content": "Tell them to ignore it."}]
+
+prompt_ids = tokenizer.apply_chat_template(
+    prompt,
+    tokenize=True,
+    add_generation_prompt=True,
+    return_tensors="pt",
+).to(model.device)
+chosen_ids = tokenizer.apply_chat_template(
+    chosen,
+    tokenize=True,
+    add_generation_prompt=False,
     return_tensors="pt",
-    padding=True,
-    truncation=True,
-    max_length=4096,
-)
-inputs = {"input_ids": encoded.to(device)} if isinstance(encoded, torch.Tensor) else {
-    key: value.to(device) for key, value in encoded.items()
-}
-
-with torch.no_grad():
-    score = model(**inputs).score.float().item()
-
-print(score)
+).to(model.device)
+rejected_ids = tokenizer.apply_chat_template(
+    rejected,
+    tokenize=True,
+    add_generation_prompt=False,
+    return_tensors="pt",
+).to(model.device)
+
+with torch.inference_mode():
+    gate = model.compute_gating(input_ids=prompt_ids)
+    chosen_score = model(
+        input_ids=chosen_ids,
+        gating_output_override=gate,
+    ).score
+    rejected_score = model(
+        input_ids=rejected_ids,
+        gating_output_override=gate,
+    ).score
+
+print({"chosen": chosen_score.item(), "rejected": rejected_score.item()})
 ```
 
-## Limitations
+For padded batches, also pass the matching `attention_mask`. Scores are intended for comparison
+within a prompt; they are not calibrated probabilities or universal utility values.
+
+## Training data
+
+The model was trained with the
+[Multi-Domain-Data-Scoring and Multi-Domain-Data-Preference-Pairs](https://github.com/mestecha/multidomain_data_scoring)
+datasets. Users are responsible for reviewing the source data and applicable terms.
+
+## Intended use and limitations
 
-This is a reward model, not a standalone chat assistant. Scores are intended for
-relative comparison and should be calibrated for each downstream use case. The
-model inherits limitations from its base model and from the annotation coverage
-of the multi-domain datasets, especially for cultural contexts not represented
-in the evaluation data.
+This checkpoint is intended for research on preference ranking, reward modeling, reranking, and
+alignment evaluation. Performance differs by domain, language, topic, and distribution. It should
+not be used as the sole decision-maker in high-impact settings, and outputs should be audited for
+biases inherited from the base model and training data.
 
-## Credits
+## License
 
-This model is based on the ArmoRM/RLHFlow reward-modeling approach and adapts it
-to custom multi-domain attributes for coherence, commonsense, empathy, and
-multicultural response quality.
+The project code is released under Apache-2.0. Use of this checkpoint is also subject to the license
+and usage conditions of the base model and training datasets.

config.json

@@ -4,6 +4,9 @@
   ],
   "attention_bias": false,
   "attention_dropout": 0.0,
+  "auto_map": {
+    "AutoModel": "modeling_custom.RewardModelWithGating"
+  },
   "bos_token_id": 128000,
   "dtype": "bfloat16",
   "eos_token_id": [
@@ -11,7 +14,35 @@
     128008,
     128009
   ],
+  "gating_active_attribute_indices": [
+    0,
+    1,
+    2,
+    3,
+    4,
+    5,
+    6,
+    7,
+    8,
+    9,
+    10,
+    11,
+    12,
+    13,
+    14,
+    15,
+    16,
+    17,
+    18,
+    19,
+    20,
+    21,
+    22
+  ],
+  "gating_dropout": 0.05,
   "gating_hidden_dim": 64,
+  "gating_learnable_logit_scale": false,
+  "gating_logit_scale": 2.0,
   "gating_n_hidden": 1,
   "gating_temperature": 2.0,
   "head_dim": 128,
@@ -43,6 +74,7 @@
     "rope_theta": 500000.0,
     "rope_type": "llama3"
   },
+  "shared_prompt_gating": true,
   "tie_word_embeddings": false,
   "transformers_version": "5.3.0",
   "use_cache": false,

results/plots/spearman_by_attribute.png → model-00001-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c2a36a5270f9be69eed9ce5bb404f24747367b9e1865fa820859dffc6c88f152
-size 119865
+oid sha256:86ca401765f707e08ad005924b409d4db9d332a95001c903a56bd1df40368657
+size 3903440936

results/plots/spearman_by_attribute_baseline.png → model-00002-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3070c0c9916a1d9d93ba3ecfe7ad627516fdb81d18644d46e204fae60aabb01e
-size 120059
+oid sha256:e7f5a022674451e32d3ff897b44ed31b8beee226a53bd9b6dae31ca3e73805ad
+size 3926025408

model.safetensors → model-00003-of-00004.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:1ce4805f114045ff00f23a13f9d7e8ac46b33ff7a7f4f345182304f7f2e3b8d4
-size 15010600834
+oid sha256:c5d03ca2426551235ee495d39c50fd0208e1d987df22ecd9617f5a36fcce68e9
+size 3926025448

model-00004-of-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1129c15e9ea55604869030314294940920965e497103974ce620e67c0477e9eb
+size 3255108794

model.safetensors.index.json

@@ -0,0 +1,305 @@
+{
+  "metadata": {
+    "total_parameters": 7505283113,
+    "total_size": 15010566226
+  },
+  "weight_map": {
+    "gating.layers.0.bias": "model-00001-of-00004.safetensors",
+    "gating.layers.0.weight": "model-00001-of-00004.safetensors",
+    "gating.layers.1.bias": "model-00001-of-00004.safetensors",
+    "gating.layers.1.weight": "model-00001-of-00004.safetensors",
+    "gating.logit_scale": "model-00001-of-00004.safetensors",
+    "model.embed_tokens.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.0.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.1.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.10.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.10.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.11.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.12.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.13.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.14.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.15.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.15.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.16.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.17.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.18.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.19.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.2.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.2.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.20.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.20.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.21.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.22.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.mlp.up_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.post_attention_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.k_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.o_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.q_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.23.self_attn.v_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.input_layernorm.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.down_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.gate_proj.weight": "model-00003-of-00004.safetensors",
+    "model.layers.24.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.24.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.24.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.24.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.24.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.24.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.25.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.26.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.27.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.28.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.29.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.3.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.3.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.30.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.30.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.input_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.down_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.gate_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.mlp.up_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.post_attention_layernorm.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.self_attn.k_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.self_attn.o_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.self_attn.q_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.31.self_attn.v_proj.weight": "model-00004-of-00004.safetensors",
+    "model.layers.4.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.4.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.mlp.up_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.post_attention_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.k_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.o_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.q_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.5.self_attn.v_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.input_layernorm.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.down_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.gate_proj.weight": "model-00001-of-00004.safetensors",
+    "model.layers.6.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.6.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.7.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.8.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.input_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.down_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.gate_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.mlp.up_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.post_attention_layernorm.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.k_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.o_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.q_proj.weight": "model-00002-of-00004.safetensors",
+    "model.layers.9.self_attn.v_proj.weight": "model-00002-of-00004.safetensors",
+    "model.norm.weight": "model-00004-of-00004.safetensors",
+    "regression_layer.weight": "model-00004-of-00004.safetensors",
+    "reward_transform_matrix": "model-00004-of-00004.safetensors"
+  }
+}

modeling_custom.py

@@ -0,0 +1,215 @@
+# modeling_custom.py
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from dataclasses import dataclass
+from typing import List, Optional
+from transformers import AutoConfig, AutoModel, PreTrainedModel
+from transformers.utils import ModelOutput
+try:
+    from .utils import TOKEN_PATTERNS_BY_MODEL_TYPE, find_token_for_gating
+except ImportError:  # Local source-tree execution.
+    from utils import TOKEN_PATTERNS_BY_MODEL_TYPE, find_token_for_gating
+
+class GatingNetwork(nn.Module):
+    def __init__(self, in_features: int, out_features: int, bias: bool = True, temperature: float = 10,
+                 logit_scale: float = 1., hidden_dim: int = 1024, n_hidden: int = 3, dropout: float = 0.0,
+                 learnable_logit_scale: bool = False,
+                 active_attribute_indices: Optional[List[int]] = None):
+        super().__init__()
+        self.temperature = temperature
+        self.logit_scale = nn.Parameter(
+            torch.ones(1) * logit_scale, requires_grad=learnable_logit_scale
+        )
+        self.dropout_prob = dropout
+        active_mask = torch.ones(out_features, dtype=torch.bool)
+        if active_attribute_indices is not None:
+            if not active_attribute_indices:
+                raise ValueError("At least one attribute must be active.")
+            if min(active_attribute_indices) < 0 or max(active_attribute_indices) >= out_features:
+                raise ValueError("active_attribute_indices contains an out-of-range index.")
+            active_mask.zero_()
+            active_mask[list(active_attribute_indices)] = True
+        # Derived from packaged config; omit from state_dict for legacy compatibility.
+        self.register_buffer("active_attribute_mask", active_mask, persistent=False)
+        layers = []
+        for _ in range(n_hidden):
+            layers.append(nn.Linear(in_features, hidden_dim))
+            in_features = hidden_dim
+        layers.append(nn.Linear(in_features, out_features, bias=bias))
+        self.layers = nn.ModuleList(layers)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        for i, layer in enumerate(self.layers):
+            x = layer(x)
+            if i < len(self.layers) - 1:
+                x = F.relu(x)
+                if self.dropout_prob > 0 and self.training:
+                    x = F.dropout(x, p=self.dropout_prob)
+        logits = x / self.temperature
+        mask = self.active_attribute_mask.to(device=logits.device)
+        logits = logits.masked_fill(~mask, torch.finfo(logits.dtype).min)
+        x = F.softmax(logits, dim=-1)
+        return x * self.logit_scale
+
+@dataclass
+class CustomOutput(ModelOutput):
+    rewards: Optional[torch.Tensor] = None
+    hidden_state: Optional[torch.Tensor] = None
+    prompt_embedding: Optional[torch.Tensor] = None
+    gating_output: Optional[torch.Tensor] = None
+    score: Optional[torch.Tensor] = None
+    logits: Optional[torch.Tensor] = None
+
+class RewardModelWithGating(PreTrainedModel):
+    """Backbone-agnostic reward model with a prompt-conditioned gating network."""
+
+    config_class = AutoConfig
+    base_model_prefix = "model"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.model = AutoModel.from_config(config)
+        config_dict = config.to_dict()
+        
+        # Default objective count for this project.
+        self.num_objectives = config_dict.get("num_objectives", 23)
+        
+        self.regression_layer = nn.Linear(config.hidden_size, self.num_objectives, bias=False)
+        self.post_init()
+        
+        # Avoid torch.eye to keep compatibility with BF16 training setups.
+        I = torch.zeros(self.num_objectives, self.num_objectives)
+        I[range(self.num_objectives), range(self.num_objectives)] = 1.
+        self.reward_transform_matrix = nn.Parameter(I)
+        self.reward_transform_matrix.requires_grad = False
+        
+        self.gating = GatingNetwork(config.hidden_size, self.num_objectives,
+                                    temperature=config_dict.get("gating_temperature", 10),
+                                    logit_scale=config_dict.get("gating_logit_scale", 1.0),
+                                    hidden_dim=config_dict.get("gating_hidden_dim", 1024),
+                                    n_hidden=config_dict.get("gating_n_hidden", 3),
+                                    dropout=config_dict.get("gating_dropout", 0.0),
+                                    learnable_logit_scale=config_dict.get("gating_learnable_logit_scale", False),
+                                    active_attribute_indices=config_dict.get(
+                                        "gating_active_attribute_indices"))
+
+    def compute_gating(
+            self,
+            input_ids: torch.LongTensor,
+            attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute routing weights from a prompt-only token sequence.
+
+        Callers should render the prompt with ``add_generation_prompt=True``.
+        The final non-padding prompt token is causal and therefore cannot see
+        either candidate response.
+        """
+        outputs = self.model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            return_dict=True,
+        )
+        hidden = outputs[0]
+        if attention_mask is None:
+            positions = torch.full(
+                (hidden.shape[0],), hidden.shape[1] - 1,
+                dtype=torch.long, device=hidden.device,
+            )
+        else:
+            token_positions = torch.arange(
+                hidden.shape[1], device=hidden.device
+            ).unsqueeze(0)
+            positions = (attention_mask.long() * token_positions).argmax(dim=-1)
+        rows = torch.arange(hidden.shape[0], device=hidden.device)
+        return self.gating(hidden[rows, positions])
+
+    def forward(
+            self,
+            input_ids: Optional[torch.LongTensor] = None,
+            attention_mask: Optional[torch.Tensor] = None,
+            position_ids: Optional[torch.LongTensor] = None,
+            past_key_values: Optional[List[torch.FloatTensor]] = None,
+            inputs_embeds: Optional[torch.FloatTensor] = None,
+            labels: Optional[torch.FloatTensor] = None,
+            use_cache: Optional[bool] = None,
+            output_attentions: Optional[bool] = None,
+            output_hidden_states: Optional[bool] = None,
+            return_dict: Optional[bool] = None,
+            gating_output_override: Optional[torch.Tensor] = None,
+    ) -> CustomOutput:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        transformer_outputs = self.model(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        tokens_hidden_states = transformer_outputs[0]
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        elif inputs_embeds is not None:
+            batch_size = inputs_embeds.shape[0]
+        else:
+            raise ValueError("Either input_ids or inputs_embeds must be provided.")
+            
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+            
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                # If no pad token is found, modulo keeps ONNX-compatible indexing.
+                sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1
+                sequence_lengths = sequence_lengths % input_ids.shape[-1]
+                sequence_lengths = sequence_lengths.to(tokens_hidden_states.device)
+            else:
+                sequence_lengths = -1
+                
+        dummy_iterator = torch.arange(batch_size, device=tokens_hidden_states.device)
+        hidden_states = tokens_hidden_states[dummy_iterator, sequence_lengths]
+        assert hidden_states.shape == (batch_size, self.config.hidden_size)
+        
+        rewards = self.regression_layer(hidden_states)
+        prompt_embedding = None
+        if gating_output_override is not None:
+            gating_output = gating_output_override.to(device=rewards.device, dtype=rewards.dtype)
+            if gating_output.ndim == 1:
+                gating_output = gating_output.unsqueeze(0)
+            if gating_output.shape != rewards.shape:
+                raise ValueError(
+                    f"gating_output_override shape {tuple(gating_output.shape)} does not "
+                    f"match rewards shape {tuple(rewards.shape)}"
+                )
+        else:
+            if input_ids is None:
+                raise ValueError("input_ids is required to compute gating token positions.")
+            model_type = getattr(self.config, "model_type", None)
+            gating_token_positions = [
+                find_token_for_gating(ids.detach().cpu().tolist(), model_type) for ids in input_ids
+            ]
+            prompt_embedding = tokens_hidden_states[dummy_iterator, gating_token_positions, :]
+            gating_output = self.gating(prompt_embedding)
+        rewards_adjusted = rewards @ self.reward_transform_matrix
+        score = torch.sum(gating_output * rewards_adjusted, dim=1)
+        
+        return CustomOutput(
+            rewards=rewards,
+            hidden_state=hidden_states,
+            prompt_embedding=prompt_embedding,
+            gating_output=gating_output,
+            score=score,
+            logits=score,
+        )
+
+
+# Backward compatibility alias for existing imports/checkpoints.
+LlamaForRewardModelWithGating = RewardModelWithGating

requirements.txt

@@ -0,0 +1,4 @@
+torch>=2.6
+transformers>=4.51
+safetensors>=0.5
+accelerate>=1.0

results/eval.json

@@ -1,604 +0,0 @@
-{
-  "model": "model/multi-domain-rm-skywork-llama-3.1-8b-it-test-196",
-  "training_metadata": {
-    "base_model_path": "Skywork/Skywork-Reward-V2-Llama-3.1-8B",
-    "stage_1_weights_path": "model/regression_weights/Skywork-Reward-V2-Llama-3.1-8B_Multi-Domain-Data-Scoring_100pct.pt",
-    "stage_2_weights_path": "model/gating_network/gating_network_Skywork-Reward-V2-Llama-3.1-8B_mo_Multi-Domain-Data-Scoring_pref_Multi-Domain-Data-Preference-Pairs_ref_null_t2.0_n30000_seed0_le0.0005_we0.0_n_1_hi64_dr0.05_ba2048_co0.03_lo2.0.pt"
-  },
-  "scoring_80pct": {
-    "evaluated": 5122,
-    "skipped": 0,
-    "attributes": {
-      "co_discourse_structure": {
-        "n": 1280,
-        "mse": 0.033118,
-        "pearson": 0.766439,
-        "spearman": 0.777125
-      },
-      "co_logical_consistency": {
-        "n": 1280,
-        "mse": 0.040838,
-        "pearson": 0.709531,
-        "spearman": 0.723998
-      },
-      "co_mutual_grounding": {
-        "n": 1280,
-        "mse": 0.036121,
-        "pearson": 0.764542,
-        "spearman": 0.775175
-      },
-      "co_overall_coherence_score": {
-        "n": 1280,
-        "mse": 0.033191,
-        "pearson": 0.759132,
-        "spearman": 0.769916
-      },
-      "co_temporal_causal_coherence": {
-        "n": 1280,
-        "mse": 0.032116,
-        "pearson": 0.739292,
-        "spearman": 0.74636
-      },
-      "co_topic_coherence": {
-        "n": 1280,
-        "mse": 0.032,
-        "pearson": 0.767388,
-        "spearman": 0.769546
-      },
-      "cs_causality": {
-        "n": 1286,
-        "mse": 0.026153,
-        "pearson": 0.629517,
-        "spearman": 0.598503
-      },
-      "cs_coherence": {
-        "n": 1286,
-        "mse": 0.009247,
-        "pearson": 0.74146,
-        "spearman": 0.753485
-      },
-      "cs_consistency": {
-        "n": 1286,
-        "mse": 0.031536,
-        "pearson": 0.572519,
-        "spearman": 0.571924
-      },
-      "cs_desire": {
-        "n": 1286,
-        "mse": 0.032596,
-        "pearson": 0.660779,
-        "spearman": 0.646299
-      },
-      "cs_empathy": {
-        "n": 1286,
-        "mse": 0.017902,
-        "pearson": 0.770497,
-        "spearman": 0.762808
-      },
-      "cs_reaction": {
-        "n": 1286,
-        "mse": 0.028105,
-        "pearson": 0.765706,
-        "spearman": 0.770195
-      },
-      "em_emotional_awareness": {
-        "n": 1274,
-        "mse": 0.022979,
-        "pearson": 0.799919,
-        "spearman": 0.810258
-      },
-      "em_emotional_validation": {
-        "n": 1274,
-        "mse": 0.020944,
-        "pearson": 0.819489,
-        "spearman": 0.836191
-      },
-      "em_helpful_response": {
-        "n": 1274,
-        "mse": 0.022296,
-        "pearson": 0.765768,
-        "spearman": 0.751327
-      },
-      "em_overall_empathy_score": {
-        "n": 1274,
-        "mse": 0.017749,
-        "pearson": 0.834078,
-        "spearman": 0.847137
-      },
-      "em_perspective_taking": {
-        "n": 1274,
-        "mse": 0.018077,
-        "pearson": 0.814856,
-        "spearman": 0.830864
-      },
-      "em_supportive_engagement": {
-        "n": 1274,
-        "mse": 0.022242,
-        "pearson": 0.811304,
-        "spearman": 0.798736
-      },
-      "mu_coherence": {
-        "n": 1282,
-        "mse": 0.003183,
-        "pearson": 0.53157,
-        "spearman": 0.595166
-      },
-      "mu_cultural_specificity": {
-        "n": 1282,
-        "mse": 0.017801,
-        "pearson": 0.805912,
-        "spearman": 0.837055
-      },
-      "mu_cultural_value": {
-        "n": 1282,
-        "mse": 0.065155,
-        "pearson": 0.226475,
-        "spearman": 0.2136
-      },
-      "mu_empathy": {
-        "n": 1282,
-        "mse": 0.01003,
-        "pearson": 0.9126,
-        "spearman": 0.83685
-      },
-      "mu_naturalness": {
-        "n": 1282,
-        "mse": 0.003088,
-        "pearson": 0.583882,
-        "spearman": 0.603499
-      }
-    },
-    "domains": {
-      "coherence": {
-        "mse": 0.034564,
-        "pearson": 0.751054,
-        "spearman": 0.760353
-      },
-      "commonsense": {
-        "mse": 0.024257,
-        "pearson": 0.69008,
-        "spearman": 0.683869
-      },
-      "empathy": {
-        "mse": 0.020715,
-        "pearson": 0.807569,
-        "spearman": 0.812419
-      },
-      "multicultural": {
-        "mse": 0.019851,
-        "pearson": 0.612088,
-        "spearman": 0.617234
-      }
-    },
-    "average": {
-      "mse": 0.025064,
-      "pearson": 0.719681,
-      "spearman": 0.72287
-    },
-    "global_score": {
-      "mean": -0.084258,
-      "std": 1.821894,
-      "min": -6.84375,
-      "max": 6.9375
-    }
-  },
-  "scoring_100pct": {
-    "evaluated": 5122,
-    "skipped": 0,
-    "attributes": {
-      "co_discourse_structure": {
-        "n": 1280,
-        "mse": 0.032518,
-        "pearson": 0.770777,
-        "spearman": 0.782585
-      },
-      "co_logical_consistency": {
-        "n": 1280,
-        "mse": 0.040172,
-        "pearson": 0.715056,
-        "spearman": 0.730289
-      },
-      "co_mutual_grounding": {
-        "n": 1280,
-        "mse": 0.035946,
-        "pearson": 0.765847,
-        "spearman": 0.776893
-      },
-      "co_overall_coherence_score": {
-        "n": 1280,
-        "mse": 0.032674,
-        "pearson": 0.763504,
-        "spearman": 0.774096
-      },
-      "co_temporal_causal_coherence": {
-        "n": 1280,
-        "mse": 0.031754,
-        "pearson": 0.743045,
-        "spearman": 0.748906
-      },
-      "co_topic_coherence": {
-        "n": 1280,
-        "mse": 0.031385,
-        "pearson": 0.772437,
-        "spearman": 0.773025
-      },
-      "cs_causality": {
-        "n": 1286,
-        "mse": 0.025816,
-        "pearson": 0.635689,
-        "spearman": 0.599198
-      },
-      "cs_coherence": {
-        "n": 1286,
-        "mse": 0.00913,
-        "pearson": 0.745643,
-        "spearman": 0.753881
-      },
-      "cs_consistency": {
-        "n": 1286,
-        "mse": 0.031209,
-        "pearson": 0.578453,
-        "spearman": 0.577192
-      },
-      "cs_desire": {
-        "n": 1286,
-        "mse": 0.032062,
-        "pearson": 0.666989,
-        "spearman": 0.652277
-      },
-      "cs_empathy": {
-        "n": 1286,
-        "mse": 0.017647,
-        "pearson": 0.774784,
-        "spearman": 0.767929
-      },
-      "cs_reaction": {
-        "n": 1286,
-        "mse": 0.027656,
-        "pearson": 0.770099,
-        "spearman": 0.773719
-      },
-      "em_emotional_awareness": {
-        "n": 1274,
-        "mse": 0.022634,
-        "pearson": 0.803611,
-        "spearman": 0.812949
-      },
-      "em_emotional_validation": {
-        "n": 1274,
-        "mse": 0.020593,
-        "pearson": 0.822651,
-        "spearman": 0.838193
-      },
-      "em_helpful_response": {
-        "n": 1274,
-        "mse": 0.022002,
-        "pearson": 0.76928,
-        "spearman": 0.752839
-      },
-      "em_overall_empathy_score": {
-        "n": 1274,
-        "mse": 0.017547,
-        "pearson": 0.835748,
-        "spearman": 0.847668
-      },
-      "em_perspective_taking": {
-        "n": 1274,
-        "mse": 0.017851,
-        "pearson": 0.817472,
-        "spearman": 0.832984
-      },
-      "em_supportive_engagement": {
-        "n": 1274,
-        "mse": 0.021954,
-        "pearson": 0.814144,
-        "spearman": 0.800278
-      },
-      "mu_coherence": {
-        "n": 1282,
-        "mse": 0.003149,
-        "pearson": 0.53863,
-        "spearman": 0.598112
-      },
-      "mu_cultural_specificity": {
-        "n": 1282,
-        "mse": 0.017417,
-        "pearson": 0.811084,
-        "spearman": 0.841164
-      },
-      "mu_cultural_value": {
-        "n": 1282,
-        "mse": 0.064614,
-        "pearson": 0.240758,
-        "spearman": 0.225872
-      },
-      "mu_empathy": {
-        "n": 1282,
-        "mse": 0.009984,
-        "pearson": 0.912989,
-        "spearman": 0.837596
-      },
-      "mu_naturalness": {
-        "n": 1282,
-        "mse": 0.003063,
-        "pearson": 0.587859,
-        "spearman": 0.60883
-      }
-    },
-    "domains": {
-      "coherence": {
-        "mse": 0.034075,
-        "pearson": 0.755111,
-        "spearman": 0.764299
-      },
-      "commonsense": {
-        "mse": 0.02392,
-        "pearson": 0.695276,
-        "spearman": 0.687366
-      },
-      "empathy": {
-        "mse": 0.02043,
-        "pearson": 0.810484,
-        "spearman": 0.814152
-      },
-      "multicultural": {
-        "mse": 0.019646,
-        "pearson": 0.618264,
-        "spearman": 0.622315
-      }
-    },
-    "average": {
-      "mse": 0.024729,
-      "pearson": 0.724198,
-      "spearman": 0.726368
-    },
-    "global_score": {
-      "mean": -0.115085,
-      "std": 1.837174,
-      "min": -7.0,
-      "max": 7.0
-    }
-  },
-  "preference": {
-    "total": 15489,
-    "correct": 14356,
-    "ties": 6,
-    "skipped": 0,
-    "accuracy": 92.6851,
-    "margin_mean": 3.905499,
-    "margin_std": 3.03703,
-    "domains": {
-      "coherence": {
-        "accuracy": 87.527,
-        "correct": 3242,
-        "total": 3704,
-        "ties": 2
-      },
-      "commonsense": {
-        "accuracy": 98.0365,
-        "correct": 4993,
-        "total": 5093,
-        "ties": 1
-      },
-      "empathy": {
-        "accuracy": 96.2669,
-        "correct": 3636,
-        "total": 3777,
-        "ties": 1
-      },
-      "multicultural": {
-        "accuracy": 85.2487,
-        "correct": 2485,
-        "total": 2915,
-        "ties": 2
-      }
-    },
-    "difficulty": {
-      "easy": {
-        "accuracy": 95.8952,
-        "correct": 11891,
-        "total": 12400,
-        "ties": 4
-      },
-      "hard": {
-        "accuracy": 73.5965,
-        "correct": 839,
-        "total": 1140,
-        "ties": 0
-      },
-      "medium": {
-        "accuracy": 83.4274,
-        "correct": 1626,
-        "total": 1949,
-        "ties": 2
-      }
-    }
-  },
-  "cultural": {
-    "evaluated": 180,
-    "skipped": 0,
-    "global_score": {
-      "mean": -0.9639,
-      "std": 1.4414,
-      "min": -6.3438,
-      "max": 3.4531
-    },
-    "countries": {
-      "Australia": {
-        "n": 15,
-        "score_mean": -0.949,
-        "score_std": 0.7752,
-        "mu_attributes": {
-          "mu_coherence": 0.875,
-          "mu_cultural_specificity": 0.5217,
-          "mu_cultural_value": 0.638,
-          "mu_empathy": 0.1304,
-          "mu_naturalness": 0.8076
-        }
-      },
-      "Chile": {
-        "n": 15,
-        "score_mean": -1.6327,
-        "score_std": 1.3299,
-        "mu_attributes": {
-          "mu_coherence": 0.876,
-          "mu_cultural_specificity": 0.5439,
-          "mu_cultural_value": 0.6354,
-          "mu_empathy": 0.1497,
-          "mu_naturalness": 0.8112
-        }
-      },
-      "China": {
-        "n": 15,
-        "score_mean": -0.7811,
-        "score_std": 1.2349,
-        "mu_attributes": {
-          "mu_coherence": 0.8802,
-          "mu_cultural_specificity": 0.5197,
-          "mu_cultural_value": 0.6831,
-          "mu_empathy": 0.1485,
-          "mu_naturalness": 0.8031
-        }
-      },
-      "Germany": {
-        "n": 15,
-        "score_mean": -1.2286,
-        "score_std": 1.4947,
-        "mu_attributes": {
-          "mu_coherence": 0.876,
-          "mu_cultural_specificity": 0.5379,
-          "mu_cultural_value": 0.6578,
-          "mu_empathy": 0.1596,
-          "mu_naturalness": 0.8005
-        }
-      },
-      "India": {
-        "n": 15,
-        "score_mean": -0.8579,
-        "score_std": 0.9938,
-        "mu_attributes": {
-          "mu_coherence": 0.8753,
-          "mu_cultural_specificity": 0.5374,
-          "mu_cultural_value": 0.6479,
-          "mu_empathy": 0.1505,
-          "mu_naturalness": 0.8042
-        }
-      },
-      "Japan": {
-        "n": 15,
-        "score_mean": -0.7905,
-        "score_std": 1.2135,
-        "mu_attributes": {
-          "mu_coherence": 0.8711,
-          "mu_cultural_specificity": 0.5289,
-          "mu_cultural_value": 0.6315,
-          "mu_empathy": 0.1584,
-          "mu_naturalness": 0.8073
-        }
-      },
-      "Kenya": {
-        "n": 15,
-        "score_mean": -0.3351,
-        "score_std": 0.9712,
-        "mu_attributes": {
-          "mu_coherence": 0.8747,
-          "mu_cultural_specificity": 0.5799,
-          "mu_cultural_value": 0.6654,
-          "mu_empathy": 0.1466,
-          "mu_naturalness": 0.8076
-        }
-      },
-      "Mexico": {
-        "n": 15,
-        "score_mean": -0.4658,
-        "score_std": 1.5211,
-        "mu_attributes": {
-          "mu_coherence": 0.8953,
-          "mu_cultural_specificity": 0.5355,
-          "mu_cultural_value": 0.6893,
-          "mu_empathy": 0.1547,
-          "mu_naturalness": 0.8031
-        }
-      },
-      "Poland": {
-        "n": 15,
-        "score_mean": -1.5578,
-        "score_std": 2.4427,
-        "mu_attributes": {
-          "mu_coherence": 0.8745,
-          "mu_cultural_specificity": 0.5493,
-          "mu_cultural_value": 0.6445,
-          "mu_empathy": 0.1511,
-          "mu_naturalness": 0.8065
-        }
-      },
-      "South Africa": {
-        "n": 15,
-        "score_mean": -0.7571,
-        "score_std": 1.4231,
-        "mu_attributes": {
-          "mu_coherence": 0.8753,
-          "mu_cultural_specificity": 0.5734,
-          "mu_cultural_value": 0.6654,
-          "mu_empathy": 0.184,
-          "mu_naturalness": 0.8112
-        }
-      },
-      "Spain": {
-        "n": 15,
-        "score_mean": -1.1197,
-        "score_std": 1.1839,
-        "mu_attributes": {
-          "mu_coherence": 0.8758,
-          "mu_cultural_specificity": 0.5292,
-          "mu_cultural_value": 0.6508,
-          "mu_empathy": 0.1525,
-          "mu_naturalness": 0.8068
-        }
-      },
-      "United States of America": {
-        "n": 15,
-        "score_mean": -1.0909,
-        "score_std": 1.4211,
-        "mu_attributes": {
-          "mu_coherence": 0.8852,
-          "mu_cultural_specificity": 0.5185,
-          "mu_cultural_value": 0.6552,
-          "mu_empathy": 0.147,
-          "mu_naturalness": 0.8185
-        }
-      }
-    },
-    "arousal": {
-      "1": {
-        "n": 12,
-        "mean": -1.1308,
-        "std": 1.2953
-      },
-      "2": {
-        "n": 32,
-        "mean": -1.301,
-        "std": 1.2291
-      },
-      "3": {
-        "n": 53,
-        "mean": -0.6617,
-        "std": 1.5538
-      },
-      "4": {
-        "n": 51,
-        "mean": -1.022,
-        "std": 1.4584
-      },
-      "5": {
-        "n": 32,
-        "mean": -0.9719,
-        "std": 1.37
-      }
-    },
-    "score_vs_arousal": {
-      "pearson": 0.0378,
-      "spearman": 0.0356
-    }
-  }
-}

results/eval_baseline.json

@@ -1,333 +0,0 @@
-{
-  "model": "Skywork/Skywork-Reward-V2-Llama-3.1-8B",
-  "type": "baseline",
-  "scoring": {
-    "evaluated": 5122,
-    "skipped": 0,
-    "attributes": {
-      "co_discourse_structure": {
-        "n": 1280,
-        "mse": 267.768484,
-        "pearson": 0.399503,
-        "spearman": 0.402577
-      },
-      "co_logical_consistency": {
-        "n": 1280,
-        "mse": 269.287788,
-        "pearson": 0.334027,
-        "spearman": 0.345432
-      },
-      "co_mutual_grounding": {
-        "n": 1280,
-        "mse": 267.451739,
-        "pearson": 0.415589,
-        "spearman": 0.419526
-      },
-      "co_overall_coherence_score": {
-        "n": 1280,
-        "mse": 268.872669,
-        "pearson": 0.389374,
-        "spearman": 0.396514
-      },
-      "co_temporal_causal_coherence": {
-        "n": 1280,
-        "mse": 270.468984,
-        "pearson": 0.366464,
-        "spearman": 0.383941
-      },
-      "co_topic_coherence": {
-        "n": 1280,
-        "mse": 270.248653,
-        "pearson": 0.373761,
-        "spearman": 0.38069
-      },
-      "cs_causality": {
-        "n": 1286,
-        "mse": 368.376983,
-        "pearson": 0.119173,
-        "spearman": 0.107379
-      },
-      "cs_coherence": {
-        "n": 1286,
-        "mse": 396.305863,
-        "pearson": 0.403845,
-        "spearman": 0.400605
-      },
-      "cs_consistency": {
-        "n": 1286,
-        "mse": 375.106672,
-        "pearson": 0.083197,
-        "spearman": 0.077648
-      },
-      "cs_desire": {
-        "n": 1286,
-        "mse": 369.624386,
-        "pearson": 0.127067,
-        "spearman": 0.120142
-      },
-      "cs_empathy": {
-        "n": 1286,
-        "mse": 385.477148,
-        "pearson": 0.222875,
-        "spearman": 0.208758
-      },
-      "cs_reaction": {
-        "n": 1286,
-        "mse": 375.427491,
-        "pearson": 0.199355,
-        "spearman": 0.188804
-      },
-      "em_emotional_awareness": {
-        "n": 1274,
-        "mse": 428.055558,
-        "pearson": 0.587765,
-        "spearman": 0.609782
-      },
-      "em_emotional_validation": {
-        "n": 1274,
-        "mse": 425.000385,
-        "pearson": 0.576845,
-        "spearman": 0.646894
-      },
-      "em_helpful_response": {
-        "n": 1274,
-        "mse": 425.850433,
-        "pearson": 0.463525,
-        "spearman": 0.452144
-      },
-      "em_overall_empathy_score": {
-        "n": 1274,
-        "mse": 426.318754,
-        "pearson": 0.618128,
-        "spearman": 0.652989
-      },
-      "em_perspective_taking": {
-        "n": 1274,
-        "mse": 426.829288,
-        "pearson": 0.620675,
-        "spearman": 0.639341
-      },
-      "em_supportive_engagement": {
-        "n": 1274,
-        "mse": 425.657313,
-        "pearson": 0.548979,
-        "spearman": 0.57158
-      },
-      "mu_coherence": {
-        "n": 1282,
-        "mse": 68.11454,
-        "pearson": 0.301372,
-        "spearman": 0.327568
-      },
-      "mu_cultural_specificity": {
-        "n": 1282,
-        "mse": 66.841459,
-        "pearson": 0.224928,
-        "spearman": 0.243449
-      },
-      "mu_cultural_value": {
-        "n": 1282,
-        "mse": 66.84707,
-        "pearson": 0.014693,
-        "spearman": 0.017508
-      },
-      "mu_empathy": {
-        "n": 1282,
-        "mse": 63.197375,
-        "pearson": 0.447345,
-        "spearman": 0.398579
-      },
-      "mu_naturalness": {
-        "n": 1282,
-        "mse": 67.880863,
-        "pearson": 0.22854,
-        "spearman": 0.227955
-      }
-    },
-    "domains": {
-      "coherence": {
-        "mse": 269.016386,
-        "pearson": 0.379786,
-        "spearman": 0.388113
-      },
-      "commonsense": {
-        "mse": 378.386424,
-        "pearson": 0.192585,
-        "spearman": 0.183889
-      },
-      "empathy": {
-        "mse": 426.285289,
-        "pearson": 0.569319,
-        "spearman": 0.595455
-      },
-      "multicultural": {
-        "mse": 66.576262,
-        "pearson": 0.243376,
-        "spearman": 0.243012
-      }
-    },
-    "average": {
-      "mse": 294.565648,
-      "pearson": 0.35074,
-      "spearman": 0.357383
-    }
-  },
-  "preference": {
-    "total": 16091,
-    "correct": 12528,
-    "ties": 46,
-    "skipped": 0,
-    "accuracy": 77.8572,
-    "margin_mean": 14.479973,
-    "margin_std": 16.697453,
-    "domains": {
-      "coherence": {
-        "accuracy": 68.363,
-        "correct": 2539,
-        "total": 3714,
-        "ties": 13
-      },
-      "commonsense": {
-        "accuracy": 98.2532,
-        "correct": 5006,
-        "total": 5095,
-        "ties": 2
-      },
-      "empathy": {
-        "accuracy": 92.0572,
-        "correct": 3477,
-        "total": 3777,
-        "ties": 0
-      },
-      "multicultural": {
-        "accuracy": 42.9672,
-        "correct": 1506,
-        "total": 3505,
-        "ties": 31
-      }
-    },
-    "difficulty": {
-      "easy": {
-        "accuracy": 86.5554,
-        "correct": 10790,
-        "total": 12466,
-        "ties": 14
-      },
-      "hard": {
-        "accuracy": 44.069,
-        "correct": 639,
-        "total": 1450,
-        "ties": 8
-      },
-      "medium": {
-        "accuracy": 50.5287,
-        "correct": 1099,
-        "total": 2175,
-        "ties": 24
-      }
-    }
-  },
-  "cultural": {
-    "evaluated": 180,
-    "skipped": 0,
-    "global_score": {
-      "mean": -4.8502,
-      "std": 6.9709,
-      "min": -24.25,
-      "max": 11.8125
-    },
-    "countries": {
-      "Australia": {
-        "n": 15,
-        "score_mean": -3.379,
-        "score_std": 7.5048
-      },
-      "Chile": {
-        "n": 15,
-        "score_mean": -2.7193,
-        "score_std": 7.3955
-      },
-      "China": {
-        "n": 15,
-        "score_mean": -8.7948,
-        "score_std": 7.0204
-      },
-      "Germany": {
-        "n": 15,
-        "score_mean": -2.6734,
-        "score_std": 7.6305
-      },
-      "India": {
-        "n": 15,
-        "score_mean": -3.3963,
-        "score_std": 6.6906
-      },
-      "Japan": {
-        "n": 15,
-        "score_mean": -6.0005,
-        "score_std": 5.6984
-      },
-      "Kenya": {
-        "n": 15,
-        "score_mean": -5.2865,
-        "score_std": 6.4503
-      },
-      "Mexico": {
-        "n": 15,
-        "score_mean": -4.4237,
-        "score_std": 6.8976
-      },
-      "Poland": {
-        "n": 15,
-        "score_mean": -4.7628,
-        "score_std": 5.3618
-      },
-      "South Africa": {
-        "n": 15,
-        "score_mean": -5.309,
-        "score_std": 6.4878
-      },
-      "Spain": {
-        "n": 15,
-        "score_mean": -7.5237,
-        "score_std": 6.7529
-      },
-      "United States of America": {
-        "n": 15,
-        "score_mean": -3.9339,
-        "score_std": 6.5357
-      }
-    },
-    "arousal": {
-      "1": {
-        "n": 12,
-        "mean": -7.4241,
-        "std": 6.301
-      },
-      "2": {
-        "n": 32,
-        "mean": -6.8282,
-        "std": 6.8316
-      },
-      "3": {
-        "n": 53,
-        "mean": -5.1853,
-        "std": 7.1044
-      },
-      "4": {
-        "n": 51,
-        "mean": -4.427,
-        "std": 6.6602
-      },
-      "5": {
-        "n": 32,
-        "mean": -2.0265,
-        "std": 6.5304
-      }
-    },
-    "score_vs_arousal": {
-      "pearson": 0.2261,
-      "spearman": 0.2092
-    }
-  }
-}

training_metadata.json

@@ -1,5 +0,0 @@
-{
-  "base_model_path": "Skywork/Skywork-Reward-V2-Llama-3.1-8B",
-  "stage_1_weights_path": "model/regression_weights/Skywork-Reward-V2-Llama-3.1-8B_Multi-Domain-Data-Scoring_100pct.pt",
-  "stage_2_weights_path": "model/gating_network/gating_network_Skywork-Reward-V2-Llama-3.1-8B_mo_Multi-Domain-Data-Scoring_pref_Multi-Domain-Data-Preference-Pairs_ref_null_t2.0_n30000_seed0_le0.0005_we0.0_n_1_hi64_dr0.05_ba2048_co0.03_lo2.0.pt"
-}

utils.py

@@ -0,0 +1,286 @@
+# utils.py — Shared utility functions for the multidomain_model pipeline.
+
+import json
+import os
+from typing import Optional, Sequence
+
+import torch
+from transformers import AutoTokenizer
+
+
+# ---------------------------------------------------------------------------
+# Remote-code detection
+# ---------------------------------------------------------------------------
+
+def _requires_remote_code(model_path: str) -> bool:
+    """Return True when the model needs trust_remote_code=True."""
+    model_path_l = str(model_path).lower()
+    return "qwen3" in model_path_l
+
+
+# ---------------------------------------------------------------------------
+# Tokenizer loading
+# ---------------------------------------------------------------------------
+
+def _load_tokenizer_robust(model_path: str):
+    """Load tokenizer with fallback to slow tokenizer when fast conversion deps are missing."""
+    trust_remote_code = _requires_remote_code(model_path)
+    try:
+        return AutoTokenizer.from_pretrained(model_path, trust_remote_code=trust_remote_code)
+    except (ValueError, ImportError) as e:
+        print(f"Warning: Fast tokenizer load failed ({e}). Retrying with use_fast=False...")
+        return AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=trust_remote_code)
+
+
+# ---------------------------------------------------------------------------
+# Dataset / file resolution
+# ---------------------------------------------------------------------------
+
+def _resolve_local_dataset_file(dataset_path: str):
+    """Resolve local JSON/JSONL path, accepting optional missing extension."""
+    candidate_paths = [dataset_path]
+    if not dataset_path.endswith(".jsonl") and not dataset_path.endswith(".json"):
+        candidate_paths.extend([f"{dataset_path}.jsonl", f"{dataset_path}.json"])
+
+    for candidate in candidate_paths:
+        if os.path.isfile(candidate):
+            return candidate
+    return None
+
+
+def _resolve_jsonl_path(path: str) -> str:
+    """Return *path* if it exists, otherwise try appending .jsonl."""
+    if os.path.isfile(path):
+        return path
+    candidate = path + ".jsonl"
+    if os.path.isfile(candidate):
+        return candidate
+    raise FileNotFoundError(f"Dataset not found: {path} (also tried {candidate})")
+
+
+def load_cultural_test(data_dir: str) -> list[dict]:
+    """Load all JSON/JSONL cultural test files from *data_dir* and return a flat list of records."""
+    records: list[dict] = []
+    if not os.path.isdir(data_dir):
+        return records
+    for fname in sorted(os.listdir(data_dir)):
+        fpath = os.path.join(data_dir, fname)
+        if fname.endswith(".jsonl"):
+            with open(fpath, "r", encoding="utf-8") as f:
+                for line in f:
+                    line = line.strip()
+                    if line:
+                        records.append(json.loads(line))
+        elif fname.endswith(".json"):
+            with open(fpath, "r", encoding="utf-8") as f:
+                rows = json.load(f)
+            if isinstance(rows, list):
+                records.extend(rows)
+            else:
+                records.append(rows)
+    return records
+
+
+def parse_cultural_conversation(record: dict) -> list[dict]:
+    """Parse a cultural test record's conversation field into chat messages.
+
+    Maps the first speaker to 'user', the second to 'assistant', and merges
+    consecutive turns from the same speaker.
+    """
+    conv = record.get("conversation", "")
+    lines = conv.split("\n")
+    messages: list[dict] = []
+    speakers: dict[str, str] = {}
+
+    for line in lines:
+        line = line.strip()
+        if not line:
+            continue
+        idx = line.find(": ")
+        if idx <= 0:
+            continue
+        speaker_id = line[:idx]
+        text = line[idx + 2:]
+        if speaker_id not in speakers:
+            speakers[speaker_id] = "user" if len(speakers) == 0 else "assistant"
+        role = speakers[speaker_id]
+        if messages and messages[-1]["role"] == role:
+            messages[-1]["content"] += "\n" + text
+        else:
+            messages.append({"role": role, "content": text})
+    return messages
+
+
+def load_jsonl_test(path: str) -> list[dict]:
+    """Load all records whose split == 'test' from a JSONL file."""
+    path = _resolve_jsonl_path(path)
+    records: list[dict] = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            record = json.loads(line)
+            split = record.get("split") or record.get("metadata", {}).get("split")
+            if split == "test":
+                records.append(record)
+    return records
+
+
+# ---------------------------------------------------------------------------
+# Save-path construction (stages 1 & 2)
+# ---------------------------------------------------------------------------
+
+def _build_save_paths(base_data_dir: str, model_name: str, dataset_folder: str, base_file_stem: str, n_shards: int, shard_idx: int):
+    """Construct output dir and filename consistently across stages."""
+    final_dir = os.path.join(base_data_dir, "embeddings", model_name, dataset_folder)
+    os.makedirs(final_dir, exist_ok=True)
+    if n_shards > 1:
+        file_name = f"{base_file_stem}-{shard_idx:05d}-of-{n_shards:05d}.safetensors"
+    else:
+        file_name = f"{base_file_stem}.safetensors"
+    return final_dir, os.path.join(final_dir, file_name)
+
+
+# ---------------------------------------------------------------------------
+# Inference model path resolution
+# ---------------------------------------------------------------------------
+
+def _resolve_inference_model_path(
+    config: dict,
+    cli_model_path: str | None,
+    cli_model_parent_dir: str | None,
+    cli_model_name: str | None,
+) -> str:
+    if cli_model_path:
+        return cli_model_path
+
+    inference_cfg = config.get("inference", {}) if isinstance(config, dict) else {}
+    if not isinstance(inference_cfg, dict):
+        inference_cfg = {}
+
+    explicit_model_path = inference_cfg.get("model_path")
+    if explicit_model_path:
+        return str(explicit_model_path)
+
+    if cli_model_parent_dir or cli_model_name:
+        model_parent_dir = str(cli_model_parent_dir or inference_cfg.get("model_parent_dir", "model"))
+        model_name = cli_model_name or inference_cfg.get("model_name")
+        if not model_name:
+            raise ValueError("model_name must be provided via --model_name or config.yaml inference.model_name")
+        return os.path.join(model_parent_dir, str(model_name))
+
+    model_name = inference_cfg.get("model_name")
+    if not model_name:
+        raise ValueError("model_name must be provided via --model_name or config.yaml inference.model_name")
+    model_parent_dir = str(inference_cfg.get("model_parent_dir", "model"))
+    return os.path.join(model_parent_dir, str(model_name))
+
+
+# ---------------------------------------------------------------------------
+# Token patterns and gating-position lookup
+# ---------------------------------------------------------------------------
+
+# Canonical mapping uses "llama3" (stage-2 convention); "llama" is an alias
+# so that modeling_custom / stage-3 lookups also resolve correctly.
+TOKEN_PATTERNS_BY_MODEL_TYPE = {
+    # Llama3: "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"
+    "llama3": [128009, 128006, 78191, 128007, 271],
+    "llama":  [128009, 128006, 78191, 128007, 271],
+    # Gemma2: "<end_of_turn>\n<start_of_turn>model\n"
+    "gemma2": [107, 108, 106, 2516, 108],
+    # Mistral Instruct: "[/INST]" marks the start of the assistant response.
+    "mistral": [733, 28748, 16289, 28793],
+}
+
+
+def find_token_for_gating(tokens: Sequence[int], model_type: Optional[str]) -> int:
+    """Return the start index of the last model-specific token pattern.
+
+    For Qwen3/auto (and any model_type without an explicit pattern), falls back
+    to the last token position.
+    """
+    if model_type == "qwen3":
+        return max(len(tokens) - 1, 0)
+
+    token_pattern = TOKEN_PATTERNS_BY_MODEL_TYPE.get(model_type)
+    if not token_pattern:
+        return max(len(tokens) - 1, 0)
+
+    token_pattern_len = len(token_pattern)
+    search_end = len(tokens)
+    for j in range(search_end - token_pattern_len, -1, -1):
+        if list(tokens[j:j + token_pattern_len]) == token_pattern:
+            return j
+    # Fallback if exact marker pattern is not present in rendered prompt.
+    return max(len(tokens) - 1, 0)
+
+
+# ---------------------------------------------------------------------------
+# Inference scoring helper
+# ---------------------------------------------------------------------------
+
+def _tokenize_chat(tokenizer, messages, device, max_length, *, add_generation_prompt=False):
+    """Render then tokenize a chat consistently across preparation and inference."""
+    text = tokenizer.apply_chat_template(
+        messages, tokenize=False, add_generation_prompt=add_generation_prompt,
+    )
+    encoding = tokenizer(
+        text, return_tensors="pt", padding=True, truncation=True, max_length=max_length,
+    )
+    return {
+        key: value.to(device) if isinstance(value, torch.Tensor) else value
+        for key, value in encoding.items()
+    }
+
+
+@torch.no_grad()
+def _score_messages(model, tokenizer, messages, device, max_length, gating_output_override=None):
+    """Tokenize chat messages and run one model forward pass."""
+    if (
+        gating_output_override is None
+        and getattr(model.config, "shared_prompt_gating", False)
+        and messages
+        and messages[-1].get("role") == "assistant"
+        and len(messages) > 1
+    ):
+        prompt_encoding = _tokenize_chat(
+            tokenizer, messages[:-1], device, max_length,
+            add_generation_prompt=True,
+        )
+        gating_output_override = model.compute_gating(
+            input_ids=prompt_encoding["input_ids"],
+            attention_mask=prompt_encoding.get("attention_mask"),
+        )
+    encoding = _tokenize_chat(tokenizer, messages, device, max_length)
+    return model(
+        input_ids=encoding["input_ids"],
+        attention_mask=encoding.get("attention_mask"),
+        gating_output_override=gating_output_override,
+    )
+
+
+@torch.no_grad()
+def _score_pair_shared_gate(
+        model, tokenizer, prompt_messages, chosen_messages, rejected_messages,
+        device, max_length,
+):
+    """Score a preference pair with one prompt-only gate shared by both candidates."""
+    prompt_encoding = _tokenize_chat(
+        tokenizer,
+        prompt_messages,
+        device,
+        max_length,
+        add_generation_prompt=True,
+    )
+    gating_output = model.compute_gating(
+        input_ids=prompt_encoding["input_ids"],
+        attention_mask=prompt_encoding.get("attention_mask"),
+    )
+    chosen = _score_messages(
+        model, tokenizer, chosen_messages, device, max_length, gating_output,
+    )
+    rejected = _score_messages(
+        model, tokenizer, rejected_messages, device, max_length, gating_output,
+    )
+    return chosen, rejected, gating_output