[init] demo

app.py

@@ -0,0 +1,572 @@
+"""
+Gradio Side-by-Side Model Comparison Demo
+
+This creates a web interface to compare three inference modes simultaneously:
+1. Single: Regular HuggingFace model
+2. T2T: Two-stage inference (shows context + answer)
+3. C2C: Rosetta model with projectors
+
+All models are loaded at startup and respond to the same input in parallel.
+"""
+
+import os
+import sys
+import torch
+import argparse
+import gradio as gr
+from pathlib import Path
+from typing import Optional, Generator
+from queue import Queue
+from threading import Thread
+
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
+from rosetta.utils.evaluate import load_rosetta_model, load_hf_model, set_default_chat_template
+from rosetta.model.wrapper import RosettaModel
+from rosetta.baseline.multi_stage import TwoStageInference
+
+
+class ModelManager:
+    """Manages loading and inference for all three model types."""
+    
+    def __init__(
+        self,
+        single_model_name: str = "Qwen/Qwen2.5-0.5B-Instruct",
+        t2t_context_model: str = "Qwen/Qwen3-0.6B",
+        t2t_answer_model: str = "Qwen/Qwen2.5-0.5B-Instruct",
+        c2c_checkpoint_path: str = "local/checkpoints/qwen3_0.6b+qwen2.5_0.5b_Fuser",
+        device: str = "cuda"
+    ):
+        """
+        Initialize ModelManager with model configurations.
+        
+        Args:
+            single_model_name: HuggingFace model name for single mode
+            t2t_context_model: Context model for T2T mode
+            t2t_answer_model: Answer model for T2T mode
+            c2c_checkpoint_path: Path to C2C checkpoint directory
+            device: Device to use (cuda, cpu, or auto)
+        """
+        if device == "auto":
+            self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        else:
+            self.device = torch.device(device)
+        print(f"Using device: {self.device}")
+        
+        # Model configurations
+        self.single_model_name = single_model_name
+        self.t2t_context_model = t2t_context_model
+        self.t2t_answer_model = t2t_answer_model
+        self.c2c_checkpoint_path = c2c_checkpoint_path
+        
+        # T2T prompt configurations
+        self.t2t_background_prompt = "Briefly describe the most useful background to answer the question:\n\n{question}"
+        self.t2t_answer_prompt = "Based on the background, answer the question:\n\n{question}"  # Format for second round question
+        self.t2t_context_max_tokens = 512
+        self.t2t_answer_max_tokens = 512
+        
+        # Generation configuration (shared across all models)
+        # To enable sampling: set use_sampling=True and adjust temperature/top_p/top_k
+        # Current mode: Greedy decoding (do_sample=False)
+        self.use_sampling = False  # Set to True to enable sampling
+        self.temperature = 0.7     # Used when use_sampling=True
+        self.top_p = 0.8          # Used when use_sampling=True
+        self.top_k = 20           # Used when use_sampling=True
+        
+        # Initialize models
+        self.single_model = None
+        self.single_tokenizer = None
+        self.t2t_model = None
+        self.c2c_model = None
+        self.c2c_tokenizer = None
+        
+        # C2C model names (will be loaded from config)
+        self.c2c_base_model = None
+        self.c2c_teacher_model = None
+        
+        print("=" * 60)
+        print("Initializing models... This may take a few minutes.")
+        print("=" * 60)
+        
+        self._load_all_models()
+    
+    def _load_single_model(self):
+        """Load single HuggingFace model."""
+        print(f"\n[Single] Loading {self.single_model_name}...")
+        self.single_model, self.single_tokenizer = load_hf_model(
+            self.single_model_name, self.device
+        )
+        set_default_chat_template(self.single_tokenizer, self.single_model_name)
+        print("[Single] ✓ Model loaded")
+    
+    def _load_t2t_model(self):
+        """Load two-stage model."""
+        print(f"\n[T2T] Loading two-stage model...")
+        print(f"  Context: {self.t2t_context_model}")
+        print(f"  Answer: {self.t2t_answer_model}")
+        print(f"  Background prompt: {self.t2t_background_prompt}")
+        print(f"  Answer prompt: {self.t2t_answer_prompt}")
+        
+        self.t2t_model = TwoStageInference(
+            context_model_path=self.t2t_context_model,
+            answer_model_path=self.t2t_answer_model,
+            device=str(self.device),
+            background_prompt=self.t2t_background_prompt
+        )
+        print("[T2T] ✓ Model loaded")
+    
+    def _load_c2c_model(self):
+        """Load Rosetta (C2C) model."""
+        print(f"\n[C2C] Loading Rosetta model from {self.c2c_checkpoint_path}...")
+        
+        # Check if checkpoint exists
+        if not Path(self.c2c_checkpoint_path).exists():
+            raise FileNotFoundError(
+                f"C2C checkpoint not found: {self.c2c_checkpoint_path}\n"
+                "You can download the checkpoints automatically with:\n"
+                "    from huggingface_hub import snapshot_download\n"
+                "    snapshot_download(\n"
+                "        repo_id='nics-efc/C2C_Fuser',\n"
+                "        allow_patterns=['qwen3_0.6b+qwen2.5_0.5b_Fuser/*'],\n"
+                "        local_dir='local/checkpoints'\n"
+                "    )\n"
+                "Or see the project README for more details."
+            )
+        
+        # Load config
+        import yaml
+        config_path = Path(self.c2c_checkpoint_path) / "config.json"
+        if not config_path.exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+        
+        with open(config_path, "r") as f:
+            config = yaml.safe_load(f)
+        
+        # Store model names from config
+        self.c2c_base_model = config["model"]["base_model"]
+        self.c2c_teacher_model = config["model"]["teacher_model"]
+        
+        # Load Rosetta model
+        subfolder_dir = Path(self.c2c_checkpoint_path) / "final"
+        if not subfolder_dir.exists():
+            raise FileNotFoundError(f"Final checkpoint directory not found: {subfolder_dir}")
+        
+        model_config = {
+            "model_name": "Rosetta",
+            "rosetta_config": {
+                "checkpoints_dir": str(subfolder_dir),
+                "base_model": self.c2c_base_model,
+                "teacher_model": self.c2c_teacher_model,
+                "is_do_alignment": config["model"].get("is_do_alignment", False),
+                "alignment_strategy": config["model"].get("alignment_strategy", "first")
+            }
+        }
+        
+        eval_config = {"checkpoints_dir": str(subfolder_dir)}
+        
+        self.c2c_model, self.c2c_tokenizer = load_rosetta_model(
+            model_config, eval_config, self.device
+        )
+        print("[C2C] ✓ Model loaded")
+    
+    def _load_all_models(self):
+        """Load all models sequentially."""
+        try:
+            self._load_single_model()
+            self._load_t2t_model()
+            self._load_c2c_model()
+            print("\n" + "=" * 60)
+            print("✓ All models loaded successfully!")
+            print("=" * 60 + "\n")
+        except Exception as e:
+            print(f"\n✗ Error loading models: {e}")
+            raise
+    
+    def _get_generation_kwargs(self, max_new_tokens: int) -> dict:
+        """
+        Get generation kwargs with consistent settings across all models.
+        
+        Args:
+            max_new_tokens: Maximum number of new tokens to generate
+            
+        Returns:
+            Dictionary of generation parameters
+        """
+        kwargs = {
+            'max_new_tokens': max_new_tokens,
+            'do_sample': self.use_sampling
+        }
+        
+        if self.use_sampling:
+            kwargs.update({
+                'temperature': self.temperature,
+                'top_p': self.top_p,
+                'top_k': self.top_k
+            })
+        
+        return kwargs
+    
+    def generate_single(self, user_input: str) -> Generator[str, None, None]:
+        """Generate response from single model with streaming."""
+        messages = [{"role": "system", "content": ""}, {"role": "user", "content": user_input}]
+        text = self.single_tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
+        )
+        inputs = self.single_tokenizer(text, return_tensors="pt").to(self.device)
+        
+        # Setup streamer
+        streamer = TextIteratorStreamer(
+            self.single_tokenizer, 
+            skip_prompt=True, 
+            skip_special_tokens=True
+        )
+        
+        # Generation parameters
+        generation_kwargs = {
+            'input_ids': inputs.input_ids,
+            'attention_mask': inputs.attention_mask,
+            'streamer': streamer,
+            **self._get_generation_kwargs(max_new_tokens=2048)
+        }
+        
+        # Run generation in separate thread
+        thread = Thread(target=self.single_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Stream tokens
+        generated_text = ""
+        for token in streamer:
+            generated_text += token
+            yield generated_text
+    
+    def generate_t2t(self, user_input: str) -> Generator[tuple[str, str], None, None]:
+        """Generate response from T2T model with streaming (returns context, answer)."""
+        # Stage 1: Context generation
+        context_streamer = TextIteratorStreamer(
+            self.t2t_model.context_tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+        
+        prompt = self.t2t_background_prompt.format(question=user_input)
+        inputs = self.t2t_model.context_tokenizer.apply_chat_template(
+            [{"role": "user", "content": prompt}],
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            enable_thinking=False
+        ).to(self.device)
+        
+        generation_kwargs = {
+            'input_ids': inputs,
+            'streamer': context_streamer,
+            **self._get_generation_kwargs(max_new_tokens=self.t2t_context_max_tokens)
+        }
+        
+        # Generate context in thread
+        thread = Thread(target=self.t2t_model.context_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Stream context tokens
+        context_text = ""
+        for token in context_streamer:
+            context_text += token
+            yield context_text, ""
+        
+        thread.join()
+        
+        # Decode full context
+        with torch.inference_mode():
+            outputs = self.t2t_model.context_model.generate(
+                inputs, **self._get_generation_kwargs(max_new_tokens=self.t2t_context_max_tokens)
+            )
+        context = self.t2t_model.context_tokenizer.batch_decode(
+            outputs[:, inputs.shape[-1]:], skip_special_tokens=True
+        )[0]
+        
+        # Stage 2: Answer generation
+        answer_streamer = TextIteratorStreamer(
+            self.t2t_model.answer_tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+        
+        # Format the second round question
+        answer_question = self.t2t_answer_prompt.format(question=user_input)
+
+        inputs = self.t2t_model.answer_tokenizer.apply_chat_template(
+            [
+                {"role": "user", "content": prompt},
+                {"role": "assistant", "content": context},
+                {"role": "user", "content": answer_question}
+            ],
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            enable_thinking=False
+        ).to(self.device)
+        
+        generation_kwargs = {
+            'input_ids': inputs,
+            'streamer': answer_streamer,
+            **self._get_generation_kwargs(max_new_tokens=self.t2t_answer_max_tokens)
+        }
+        
+        # Generate answer in thread
+        thread = Thread(target=self.t2t_model.answer_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Stream answer tokens
+        answer_text = ""
+        for token in answer_streamer:
+            answer_text += token
+            yield context_text, answer_text
+    
+    def generate_c2c(self, user_input: str) -> Generator[str, None, None]:
+        """Generate response from C2C model with streaming."""
+        messages = [{"role": "system", "content": ""}, {"role": "user", "content": user_input}]
+        text = self.c2c_tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=False
+        )
+        inputs = self.c2c_tokenizer(text, return_tensors="pt").to(self.device)
+        
+        # Setup streamer
+        streamer = TextIteratorStreamer(
+            self.c2c_tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True
+        )
+        
+        # Prepare C2C-specific inputs
+        full_length = inputs.input_ids.shape[1]
+        instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(
+            full_length - 1, 1
+        ).unsqueeze(0).to(self.device)
+        label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(
+            1, 1
+        ).unsqueeze(0).to(self.device)
+        position_ids = inputs.attention_mask.long().cumsum(-1) - 1 if inputs.attention_mask is not None else \
+                      torch.arange(full_length, dtype=torch.long).unsqueeze(0).to(self.device)
+        
+        # Generation parameters
+        generation_kwargs = {
+            'kv_cache_index': [instruction_index, label_index],
+            'input_ids': inputs.input_ids,
+            'attention_mask': inputs.attention_mask,
+            'position_ids': position_ids,
+            'streamer': streamer,
+            **self._get_generation_kwargs(max_new_tokens=2048)
+        }
+        
+        # Run generation in separate thread
+        thread = Thread(target=self.c2c_model.generate, kwargs=generation_kwargs)
+        thread.start()
+        
+        # Stream tokens
+        generated_text = ""
+        for token in streamer:
+            generated_text += token
+            yield generated_text
+
+
+def create_demo(model_manager: ModelManager):
+    """Create Gradio interface."""
+    
+    # Preset example questions
+    EXAMPLE_QUESTIONS = {
+        "example1": """Instead of asking why the act of destroying the environment might be immoral, Hill wants to ask ...
+
+A. Why the act of destroying nature might be immoral.
+B. Why people who destroy the environment might be bad people.
+C. How the decision to preserve the environment benefits the environment.
+D. Whether plants have interests.""",
+        "example2": "Which company launched the Gemini 1.5 Pro model in early 2024?"
+    }
+    
+    def respond(user_input: str):
+        """Main response function that yields updates for all three models."""
+        if not user_input.strip():
+            yield "", "", "", ""
+        
+        # Generators for each model
+        single_gen = model_manager.generate_single(user_input)
+        t2t_gen = model_manager.generate_t2t(user_input)
+        c2c_gen = model_manager.generate_c2c(user_input)
+        
+        single_done = False
+        t2t_done = False
+        c2c_done = False
+        
+        single_text = ""
+        t2t_context = ""
+        t2t_answer = ""
+        c2c_text = ""
+        
+        # Stream from all three models
+        while not (single_done and t2t_done and c2c_done):
+            # Update single
+            if not single_done:
+                try:
+                    single_text = next(single_gen)
+                except StopIteration:
+                    single_done = True
+            
+            # Update T2T
+            if not t2t_done:
+                try:
+                    t2t_context, t2t_answer = next(t2t_gen)
+                except StopIteration:
+                    t2t_done = True
+            
+            # Update C2C
+            if not c2c_done:
+                try:
+                    c2c_text = next(c2c_gen)
+                except StopIteration:
+                    c2c_done = True
+            
+            # Yield current state
+            yield single_text, t2t_context, t2t_answer, c2c_text
+    
+    # Create Gradio interface
+    with gr.Blocks(title="C2C Demo", theme=gr.themes.Base()) as demo:
+        # Header with logo
+        with gr.Row():
+            with gr.Column(scale=1, min_width=100):
+                gr.Image("resource/logo.png", show_label=False, show_download_button=False, container=False, height=80)
+            with gr.Column(scale=5):
+                gr.Markdown("# Cache-to-Cache Communication Demo")
+                gr.Markdown("Compare three inference modes side-by-side: **Single** | **Text-to-Text Communication** | **Cache-to-Cache Communication**")
+        
+        gr.Markdown("---")
+        
+        # Input section
+        gr.Markdown("## Question")
+
+        # Preset question examples
+        gr.Markdown("Example Questions:")
+        with gr.Row():
+            example1_btn = gr.Button("📝 Example 1: Philosophy", size="sm")
+            example2_btn = gr.Button("📝 Example 2: Knowledge Cutoff", size="sm")
+
+
+        with gr.Row():
+            user_input = gr.Textbox(
+                label="",
+                placeholder="Type your question here...",
+                lines=2,
+                scale=4,
+                show_label=False
+            )
+        
+        with gr.Row():
+            submit_btn = gr.Button("🚀 Submit", variant="primary", scale=1)
+            clear_btn = gr.Button("🗑️ Clear", scale=1)
+        
+        gr.Markdown("---")
+        
+        # Output section - three columns
+        gr.Markdown("## Responses")
+        with gr.Row():
+            # Single column
+            with gr.Column():
+                gr.Markdown("### Single Model")
+                gr.Markdown(f"*{model_manager.single_model_name}*")
+                single_output = gr.Textbox(
+                    label="",
+                    lines=18,
+                    max_lines=30,
+                    interactive=False,
+                    show_label=False
+                )
+            
+            # T2T column (with two sub-boxes)
+            with gr.Column():
+                gr.Markdown("### Text-to-Text Communication")
+                gr.Markdown(f"*{model_manager.t2t_context_model} → {model_manager.t2t_answer_model}*")
+                t2t_context_output = gr.Textbox(
+                    label="📝 Context",
+                    lines=6,
+                    max_lines=12,
+                    interactive=False
+                )
+                t2t_answer_output = gr.Textbox(
+                    label="💬 Answer",
+                    lines=7,
+                    max_lines=14,
+                    interactive=False
+                )
+            
+            # C2C column
+            with gr.Column():
+                gr.Markdown("### Cache-to-Cache Communication")
+                gr.Markdown(f"*{model_manager.c2c_base_model} → {model_manager.c2c_teacher_model}*")
+                c2c_output = gr.Textbox(
+                    label="",
+                    lines=18,
+                    max_lines=30,
+                    interactive=False,
+                    show_label=False
+                )
+        
+        # Event handlers
+        submit_btn.click(
+            fn=respond,
+            inputs=[user_input],
+            outputs=[single_output, t2t_context_output, t2t_answer_output, c2c_output]
+        )
+        
+        user_input.submit(
+            fn=respond,
+            inputs=[user_input],
+            outputs=[single_output, t2t_context_output, t2t_answer_output, c2c_output]
+        )
+        
+        clear_btn.click(
+            fn=lambda: ("", "", "", "", ""),
+            inputs=None,
+            outputs=[user_input, single_output, t2t_context_output, t2t_answer_output, c2c_output]
+        )
+        
+        # Example question handlers
+        example1_btn.click(
+            fn=lambda: EXAMPLE_QUESTIONS["example1"],
+            inputs=None,
+            outputs=[user_input]
+        )
+        
+        example2_btn.click(
+            fn=lambda: EXAMPLE_QUESTIONS["example2"],
+            inputs=None,
+            outputs=[user_input]
+        )
+    
+    return demo
+
+
+def main():
+    """Main entry point."""
+    print("=" * 60)
+    print("Model Comparison Demo - Gradio Interface")
+    print("=" * 60)
+    
+    # Initialize models
+    model_manager = ModelManager()
+    
+    # Create and launch demo
+    demo = create_demo(model_manager)
+    
+    print("\n" + "=" * 60)
+    print("🚀 Launching Gradio interface...")
+    print("=" * 60)
+    
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        show_error=True
+    )
+
+
+if __name__ == "__main__":
+    main()
+

requirements.txt

@@ -0,0 +1,20 @@
+# Core dependencies (from pyproject.toml)
+torch==2.6.0
+transformers==4.52.4
+
+# Gradio web interface
+gradio==5.9.1
+
+# HuggingFace Hub for checkpoint downloads
+huggingface-hub>=0.26.0
+
+# Configuration file parsing
+pyyaml>=6.0
+
+# Additional dependencies for rosetta package
+datasets>=2.0
+accelerate>=0.20
+
+# Utilities
+einops>=0.8
+

rosetta/baseline/multi_stage.py

@@ -0,0 +1,992 @@
+"""
+Multi-stage evaluation utilities for VLM+LLM and LLM+LLM pipelines.
+
+This module provides utilities for multi-stage evaluation where:
+1. VLM describes/analyzes images + LLM performs reasoning
+2. LLM provides background context + LLM performs reasoning
+"""
+
+from typing import Dict, Optional, Any
+import torch
+from transformers import (
+    # Qwen2_5_VLForConditionalGeneration,
+    AutoProcessor,
+    AutoTokenizer,
+    AutoModelForCausalLM,
+)
+from rosetta.utils.evaluate import set_default_chat_template, apply_generation_config
+
+try:
+    from qwen_vl_utils import process_vision_info
+    from transformers import Qwen2_5_VLForConditionalGeneration
+except ImportError:
+    print("Please install qwen-vl-utils to use VLM models")
+
+class TwoStageInference:
+    """Two-stage LLM+LLM inference pipeline for question answering."""
+    
+    def __init__(
+        self,
+        context_model_path: str,
+        answer_model_path: str,
+        device: str = "cuda",
+        max_new_tokens: int = 1024,
+        background_prompt: str = "Briefly describe the most useful background to solve the problem:\n\n{question}",
+        generation_config: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Initialize two-stage LLM pipeline.
+        
+        Args:
+            context_model_path: Path to context-providing LLM
+            answer_model_path: Path to answer-generating LLM
+            device: Device to use
+            max_new_tokens: Maximum number of new tokens to generate
+            background_prompt: Prompt template for background generation
+            generation_config: Optional generation configuration to apply to models
+        """
+        self.device = device
+        self.max_new_tokens = max_new_tokens
+        self.background_prompt = background_prompt
+        self.generation_config = generation_config or {}
+        self._load_models(context_model_path, answer_model_path)
+    
+    def _load_models(self, context_path: str, answer_path: str):
+        """Load both LLM models."""
+        # Load context LLM
+        self.context_tokenizer = AutoTokenizer.from_pretrained(context_path)
+        # for gemma, set sliding_window=4096
+        if context_path == "google/gemma-3-1b-it":
+            torch._dynamo.config.cache_size_limit = 64
+            self.context_model = AutoModelForCausalLM.from_pretrained(
+                context_path, torch_dtype=torch.bfloat16, device_map={"": self.device}, sliding_window=4096
+            )
+        else:
+            self.context_model = AutoModelForCausalLM.from_pretrained(
+                context_path, torch_dtype=torch.bfloat16, device_map={"": self.device}
+            )
+        # Apply generation config to context model
+        apply_generation_config(self.context_model, self.generation_config)
+        
+        # Load answer LLM
+        self.answer_tokenizer = AutoTokenizer.from_pretrained(answer_path)
+        self.answer_model = AutoModelForCausalLM.from_pretrained(
+            answer_path, torch_dtype=torch.bfloat16, device_map={"": self.device}
+        )
+        # Apply generation config to answer model
+        apply_generation_config(self.answer_model, self.generation_config)
+    
+    def get_background_context(
+        self,
+        question: str,
+        max_new_tokens: Optional[int] = None
+    ) -> str:
+        """
+        Get background context from the first LLM.
+        
+        Args:
+            question: Question text (without options)
+            max_new_tokens: Max tokens to generate (uses instance default if None)
+            
+        Returns:
+            Background context
+        """
+        prompt = self.background_prompt.format(question=question)
+        messages = [{"role": "user", "content": prompt}]
+        
+        template_kwargs = {'enable_thinking': False}
+        
+        inputs = self.context_tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            **template_kwargs
+        )
+        inputs = inputs.to(self.device)
+        
+        if max_new_tokens is None:
+            max_new_tokens = self.max_new_tokens
+        
+        with torch.inference_mode():
+            outputs = self.context_model.generate(
+                inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False
+            )
+        
+        generated_ids = outputs[:, inputs.shape[-1]:]
+        context = self.context_tokenizer.batch_decode(
+            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        
+        return context
+    
+    def answer_with_context(
+        self,
+        question: str,
+        context: str,
+        max_new_tokens: Optional[int] = None,
+        original_question: Optional[str] = None
+    ) -> str:
+        """
+        Answer question using the second LLM with context.
+        
+        Args:
+            question: Full question with options and proper template
+            context: Background context from first LLM
+            max_new_tokens: Max tokens to generate (uses instance default if None)
+            original_question: Original question asked to first LLM (for conversation format)
+            
+        Returns:
+            Generated answer
+        """
+        # Use conversation format: user asks for background, assistant provides it, user asks main question
+        if original_question:
+            messages = [
+                {"role": "user", "content": self.background_prompt.format(question=original_question)},
+                {"role": "assistant", "content": context},
+                {"role": "user", "content": question}
+            ]
+        else:
+            # Fallback to simple format
+            messages = [{"role": "user", "content": f"Background context: {context}\n\n{question}"}]
+
+        template_kwargs = {'enable_thinking': False}
+        
+        inputs = self.answer_tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            **template_kwargs
+        )
+        inputs = inputs.to(self.device)
+        
+        if max_new_tokens is None:
+            max_new_tokens = self.max_new_tokens
+        
+        with torch.inference_mode():
+            outputs = self.answer_model.generate(
+                inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False
+            )
+        
+        generated_ids = outputs[:, inputs.shape[-1]:]
+        answer = self.answer_tokenizer.batch_decode(
+            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        
+        return answer
+    
+    def forward_with_context(
+        self,
+        question: str,
+        context: str,
+        original_question: Optional[str] = None,
+        response_text: Optional[str] = None,
+        **forward_kwargs
+    ) -> Any:
+        """
+        Run a forward pass on the answer model using provided context (logits mode).
+
+        Args:
+            question: Full question with options and proper template
+            context: Background context from first LLM
+            original_question: Original question asked to first LLM (for conversation format)
+            response_text: Optional text to append after the chat template to steer next-token logits
+            **forward_kwargs: Extra kwargs forwarded to the model's forward
+
+        Returns:
+            Model outputs from the forward pass (e.g., logits)
+        """
+        # Use conversation format: user asks for background, assistant provides it, user asks main question
+        if original_question:
+            messages = [
+                {"role": "user", "content": self.background_prompt.format(question=original_question)},
+                {"role": "assistant", "content": context},
+                {"role": "user", "content": question}
+            ]
+        else:
+            # Fallback to simple format
+            messages = [{"role": "user", "content": f"Background context: {context}\n\n{question}"}]
+
+        template_kwargs = {'enable_thinking': False}
+
+        # Build model inputs; if response_text is provided, append it to steer next-token prediction
+        if response_text is not None:
+            # Build raw text then append response_text
+            text = self.answer_tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                **template_kwargs
+            )
+            text = text + response_text
+            tokenized = self.answer_tokenizer(text, return_tensors="pt")
+        else:
+            # Directly build tensors with generation prompt (predict next assistant token)
+            tokenized = self.answer_tokenizer.apply_chat_template(
+                messages,
+                tokenize=True,
+                add_generation_prompt=True,
+                return_tensors="pt",
+                **template_kwargs
+            )
+
+        inputs = {k: v.to(self.device) for k, v in tokenized.items()}
+
+        with torch.inference_mode():
+            outputs = self.answer_model(**inputs, **forward_kwargs)
+
+        return outputs
+
+    def forward(
+        self,
+        question_without_options: str,
+        question_with_options: str,
+        response_text: Optional[str] = None,
+        **forward_kwargs
+    ) -> Any:
+        """
+        Two-stage forward pass (logits mode):
+        1) Generate background context with the context model
+        2) Run a forward pass on the answer model conditioned on that context
+
+        Args:
+            question_without_options: Question text without multiple choice options
+            question_with_options: Full question with options and proper template
+            response_text: Optional text appended after the chat template to steer next-token logits
+            **forward_kwargs: Extra kwargs forwarded to the model's forward
+
+        Returns:
+            Model outputs from the forward pass (e.g., logits)
+        """
+        context = self.get_background_context(question_without_options)
+        return self.forward_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options,
+            response_text=response_text,
+            **forward_kwargs
+        )
+
+    def logits_with_context(
+        self,
+        question_without_options: str,
+        question_with_options: str,
+        response_text: Optional[str] = None,
+        **forward_kwargs
+    ) -> Any:
+        """
+        Two-stage logits helper that also returns the generated background context
+        for logging as CoT.
+
+        Returns:
+            (outputs, context)
+        """
+        context = self.get_background_context(question_without_options)
+        outputs = self.forward_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options,
+            response_text=response_text,
+            **forward_kwargs
+        )
+        return outputs, context
+
+    def generate(
+        self,
+        question_without_options: str,
+        question_with_options: str,
+        communication_max_new_tokens: Optional[int] = None,
+        response_max_new_tokens: Optional[int] = None,
+        **kwargs
+    ) -> str:
+        """
+        Generate answer using two-stage processing.
+        
+        Args:
+            question_without_options: Question text without multiple choice options
+            question_with_options: Full question with options and proper template
+            communication_max_new_tokens: Maximum tokens to generate for the background context
+            response_max_new_tokens: Maximum tokens to generate for the answer
+            **kwargs: Additional generation parameters (ignored for compatibility)
+            
+        Returns:
+            Generated answer string
+        """
+        # Stage 1: Get background context
+        context = self.get_background_context(question_without_options, communication_max_new_tokens)
+        
+        # Stage 2: Answer question with context
+        answer = self.answer_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options,
+            max_new_tokens=response_max_new_tokens
+        )
+        
+        return answer
+    
+    def process(
+        self,
+        question_without_options: str,
+        question_with_options: str
+    ) -> Dict[str, str]:
+        """
+        Full two-stage processing (legacy method for backward compatibility).
+        
+        Args:
+            question_without_options: Question text without multiple choice options
+            question_with_options: Full question with options and proper template
+            
+        Returns:
+            Dictionary with context and answer
+        """
+        # Stage 1: Get background context
+        context = self.get_background_context(question_without_options)
+        
+        # Stage 2: Answer question with context
+        answer = self.answer_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options
+        )
+        
+        return {
+            "context": context,
+            "answer": answer
+        }
+
+
+class TwoStageRosetta(TwoStageInference):
+    """Two-stage LLM+Rosetta inference pipeline for question answering."""
+    
+    def __init__(
+        self,
+        context_model_path: str,
+        rosetta_checkpoint_dir: str,
+        rosetta_subfolder: str = "final",
+        device: str = "cuda",
+        max_new_tokens: int = 1024,
+        background_prompt: str = "Briefly describe the most useful background to solve the problem:\n\n{question}",
+        generation_config: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Initialize two-stage pipeline with Rosetta as second model.
+        
+        Args:
+            context_model_path: Path to context-providing LLM
+            rosetta_checkpoint_dir: Path to Rosetta checkpoint directory
+            rosetta_subfolder: Subfolder name in checkpoint directory (e.g., 'final', 'checkpoint-1000')
+            device: Device to use
+            max_new_tokens: Maximum number of new tokens to generate
+            background_prompt: Prompt template for background generation
+            generation_config: Optional generation configuration to apply to models
+        """
+        # Initialize parent class with dummy answer model path
+        # We'll override the answer model loading
+        super().__init__(
+            context_model_path=context_model_path,
+            answer_model_path=None,  # Will be overridden
+            device=device,
+            max_new_tokens=max_new_tokens,
+            background_prompt=background_prompt,
+            generation_config=generation_config
+        )
+        
+        self.rosetta_checkpoint_dir = rosetta_checkpoint_dir
+        self.rosetta_subfolder = rosetta_subfolder
+        self._load_rosetta_model()
+    
+    def _load_models(self, context_path: str, answer_path: str):
+        """
+        Override parent class _load_models to prevent loading dummy answer model.
+        We only load the context model here, and the Rosetta model is loaded separately.
+        """
+        # Only load context LLM (answer model is replaced by Rosetta)
+        self.context_tokenizer = AutoTokenizer.from_pretrained(context_path)
+        self.context_model = AutoModelForCausalLM.from_pretrained(
+            context_path, torch_dtype=torch.bfloat16, device_map={"": self.device}
+        )
+        # Apply generation config to context model
+        apply_generation_config(self.context_model, self.generation_config)
+        
+        # Skip loading answer model - we use Rosetta instead
+        print(f"Loaded context model from {context_path}")
+        print("Skipping answer model loading - using Rosetta model instead")
+    
+    def _load_rosetta_model(self):
+        """Load Rosetta model and related components following load_model_from_checkpoint pattern."""
+        import json
+        from pathlib import Path
+        from rosetta.utils.evaluate import load_rosetta_model
+        
+        checkpoint_path = Path(self.rosetta_checkpoint_dir)
+        
+        # Load config
+        config_path = checkpoint_path / "config.json"
+        if not config_path.exists():
+            raise FileNotFoundError(f"Config file not found: {config_path}")
+        
+        with open(config_path, 'r') as f:
+            config = json.load(f)
+        
+        # Check if this is a Rosetta model (has projectors)
+        subfolder_dir = checkpoint_path / self.rosetta_subfolder
+        has_projectors = subfolder_dir.exists() and any(
+            f.name.startswith("projector_") and f.name.endswith(".pt") 
+            for f in subfolder_dir.iterdir()
+        )
+        
+        if not has_projectors:
+            raise ValueError(f"No projectors found in {subfolder_dir}. This doesn't appear to be a Rosetta checkpoint.")
+        
+        # Load Rosetta model (following load_model_from_checkpoint pattern)
+        print(f"Loading Rosetta model from {self.rosetta_checkpoint_dir}")
+        
+        # Create model config for Rosetta loading
+        model_config = {
+            "model_name": "Rosetta",
+            "rosetta_config": {
+                "checkpoints_dir": str(subfolder_dir),
+                "base_model": config["model"]["base_model"],
+                "teacher_model": config["model"]["teacher_model"],
+                "is_do_alignment": config["model"].get("is_do_alignment", False),
+                "alignment_strategy": config["model"].get("alignment_strategy", "first")
+            }
+        }
+
+        print(f"Model config: {model_config}")
+        
+        eval_config = {
+            "checkpoints_dir": str(subfolder_dir)
+        }
+        
+        # Load Rosetta model using the existing utility
+        self.rosetta_model, self.rosetta_tokenizer = load_rosetta_model(
+            model_config, 
+            eval_config, 
+            device=self.device
+        )
+        
+        # Load LLM tokenizer for alignment if needed
+        is_do_alignment = config["model"].get("is_do_alignment", False)
+        llm_model_path = config["model"].get("teacher_model")
+        self.llm_tokenizer = None
+        
+        if is_do_alignment and llm_model_path:
+            try:
+                self.llm_tokenizer = AutoTokenizer.from_pretrained(str(llm_model_path))
+                if self.llm_tokenizer.pad_token is None:
+                    self.llm_tokenizer.pad_token = self.llm_tokenizer.eos_token
+                set_default_chat_template(self.llm_tokenizer, llm_model_path)
+            except Exception as e:
+                print(f"Failed to load LLM tokenizer '{llm_model_path}': {e}")
+                self.llm_tokenizer = None
+        
+        print(f"Initialized TwoStageRosetta with Rosetta model on {self.device}")
+    
+    def _prepare_rosetta_inputs(
+        self,
+        question: str,
+        context: str,
+        original_question: Optional[str] = None,
+        max_new_tokens: Optional[int] = None,
+        answer_method: str = "generate",
+        response_text: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Prepare inputs for Rosetta model using the simpler approach from live_chat_example.py.
+        
+        Args:
+            question: Question to answer
+            context: Background context from first LLM
+            original_question: Original question asked to first LLM (for conversation format)
+            max_new_tokens: Max tokens to generate (uses instance default if None)
+            
+        Returns:
+            Dictionary with prepared inputs for Rosetta model
+        """
+        # Use conversation format: user asks for background, assistant provides it, user asks main question
+        if original_question:
+            messages = [
+                {"role": "user", "content": self.background_prompt.format(question=original_question)},
+                {"role": "assistant", "content": context},
+                {"role": "user", "content": question}
+            ]
+        else:
+            # Fallback to simple format
+            messages = [{"role": "user", "content": f"Background context: {context}\n\n{question}"}]
+
+        # Apply chat template (following live_chat_example.py pattern)
+        base_text = None
+        if hasattr(self.rosetta_tokenizer, 'apply_chat_template'):
+            base_text = self.rosetta_tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=False
+            )
+        else:
+            base_text = f"### Human: {question}\n### Assistant:"
+
+        # Optionally append response_text for logits mode to steer next-token logits
+        if answer_method == 'logits' and response_text is not None:
+            text = base_text + response_text
+        else:
+            text = base_text
+
+        # Tokenize input
+        inputs = self.rosetta_tokenizer(text, return_tensors="pt").to(self.device)
+        
+        # Create kv_cache_index for Rosetta model
+        full_length = inputs.input_ids.shape[1]
+        if answer_method == 'logits':
+            # Compute response length as the extra tokens appended by response_text
+            if response_text is not None:
+                base_tok = self.rosetta_tokenizer(base_text, return_tensors="pt")
+                response_length = int(inputs.input_ids.shape[1] - base_tok.input_ids.shape[1])
+                response_length = max(response_length, 0)
+            else:
+                response_length = 0
+            instr_len = max(full_length - response_length, 0)
+            instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(instr_len, 1).unsqueeze(0).to(self.device)
+            if response_length > 0:
+                response_index = torch.tensor([-1, 0], dtype=torch.long).repeat(response_length, 1).unsqueeze(0).to(self.device)
+                kv_cache_list = [instruction_index, response_index]
+            else:
+                kv_cache_list = [instruction_index]
+        else:
+            # Generate: treat the last position as response (length 1)
+            instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(full_length - 1, 1).unsqueeze(0).to(self.device)
+            label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(1, 1).unsqueeze(0).to(self.device)
+            kv_cache_list = [instruction_index, label_index]
+        
+        # Add position_ids if needed
+        if inputs.attention_mask is None:
+            position_ids = torch.arange(inputs.input_ids.shape[-1], dtype=torch.long).unsqueeze(0).to(self.device)
+        else:
+            position_ids = inputs.attention_mask.long().cumsum(-1) - 1
+        
+        outputs = {
+            "inputs": {
+                "input_ids": inputs.input_ids,
+                "attention_mask": inputs.attention_mask,
+                "position_ids": position_ids,
+                "kv_cache_index": kv_cache_list
+            },
+            "printable_text": text
+        }
+
+        return outputs
+    
+    def answer_with_context(
+        self,
+        question: str,
+        context: str,
+        max_new_tokens: Optional[int] = None,
+        original_question: Optional[str] = None
+    ) -> str:
+        """
+        Answer question using Rosetta model with context.
+        Overrides parent class method to use Rosetta instead of regular LLM.
+        
+        Args:
+            question: Question to answer
+            context: Background context from first LLM
+            max_new_tokens: Max tokens to generate (uses instance default if None)
+            original_question: Original question asked to first LLM (for conversation format)
+            
+        Returns:
+            Generated answer
+        """
+        # Prepare inputs for Rosetta model
+        prepared = self._prepare_rosetta_inputs(
+            question=question,
+            context=context,
+            original_question=original_question,
+            max_new_tokens=max_new_tokens
+        )
+        
+        if max_new_tokens is None:
+            max_new_tokens = self.max_new_tokens
+        
+        # Generation parameters (following live_chat_example.py pattern)
+        sampling_params = {
+            'do_sample': False,
+            'max_new_tokens': max_new_tokens
+        }
+        
+        # Generate using Rosetta model (following live_chat_example.py pattern)
+        input_length = prepared['inputs']['input_ids'].shape[1]
+
+        with torch.inference_mode():
+            outputs = self.rosetta_model.generate(
+                kv_cache_index=prepared['inputs']['kv_cache_index'],
+                input_ids=prepared['inputs']['input_ids'],
+                attention_mask=prepared['inputs']['attention_mask'],
+                position_ids=prepared['inputs']['position_ids'],
+                **sampling_params
+            )
+            generated_ids = outputs[0]
+        
+        # Decode response
+        answer = self.rosetta_tokenizer.decode(generated_ids[input_length:], skip_special_tokens=True).strip()
+        
+        return answer
+
+    def forward_with_context(
+        self,
+        question: str,
+        context: str,
+        original_question: Optional[str] = None,
+        response_text: Optional[str] = None,
+        **forward_kwargs
+    ) -> Any:
+        """
+        Run a forward pass on the Rosetta model using provided context (logits mode).
+
+        Args:
+            question: Full question with options and proper template
+            context: Background context from first LLM
+            original_question: Original question asked to first LLM (for conversation format)
+            response_text: Optional text appended after the chat template to steer next-token logits
+            **forward_kwargs: Extra kwargs forwarded to the model's forward
+
+        Returns:
+            Model outputs from the forward pass (e.g., logits)
+        """
+        prepared = self._prepare_rosetta_inputs(
+            question=question,
+            context=context,
+            original_question=original_question,
+            answer_method='logits',
+            response_text=response_text
+        )
+
+        inputs = prepared['inputs']
+        with torch.inference_mode():
+            outputs = self.rosetta_model.forward(
+                kv_cache_index=inputs['kv_cache_index'],
+                input_ids=inputs['input_ids'],
+                attention_mask=inputs['attention_mask'],
+                position_ids=inputs['position_ids'],
+                **forward_kwargs
+            )
+        return outputs
+
+    def forward(
+        self,
+        question_without_options: str,
+        question_with_options: str,
+        response_text: Optional[str] = None,
+        **forward_kwargs
+    ) -> Any:
+        """
+        Two-stage forward pass (logits mode) for Rosetta:
+        1) Generate background context with the context model
+        2) Run a forward pass on the Rosetta model conditioned on that context
+
+        Args:
+            question_without_options: Question text without multiple choice options
+            question_with_options: Full question with options and proper template
+            response_text: Optional text appended after the chat template to steer next-token logits
+            **forward_kwargs: Extra kwargs forwarded to the model's forward
+
+        Returns:
+            Model outputs from the forward pass (e.g., logits)
+        """
+        # Work in progress
+        raise NotImplementedError
+        context = self.get_background_context(question_without_options)
+        return self.forward_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options,
+            response_text=response_text,
+            **forward_kwargs
+        )
+
+    def logits_with_context(
+        self,
+        question_without_options: str,
+        question_with_options: str,
+        response_text: Optional[str] = None,
+        **forward_kwargs
+    ) -> Any:
+        """
+        Two-stage logits helper that also returns the generated background context
+        for logging as CoT (Rosetta backend).
+
+        Returns:
+            (outputs, context)
+        """
+        context = self.get_background_context(question_without_options)
+        outputs = self.forward_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options,
+            response_text=response_text,
+            **forward_kwargs
+        )
+        return outputs, context
+    
+    def generate(
+        self,
+        question_without_options: str,
+        question_with_options: str,
+        max_new_tokens: Optional[int] = None,
+        **kwargs
+    ) -> str:
+        """
+        Generate answer using two-stage processing with Rosetta.
+        
+        Args:
+            question_without_options: Question text without multiple choice options
+            question_with_options: Full question with options and proper template
+            max_new_tokens: Maximum tokens to generate (passed to both stages)
+            **kwargs: Additional generation parameters (ignored for compatibility)
+            
+        Returns:
+            Generated answer string
+        """
+        # Stage 1: Get background context (uses parent class method)
+        context = self.get_background_context(question_without_options, max_new_tokens)
+        
+        # Stage 2: Answer question with context using Rosetta
+        answer = self.answer_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options,
+            max_new_tokens=max_new_tokens
+        )
+        
+        return answer
+    
+    def process(
+        self,
+        question_without_options: str,
+        question_with_options: str
+    ) -> Dict[str, str]:
+        """
+        Full two-stage processing with Rosetta (legacy method for backward compatibility).
+        
+        Args:
+            question_without_options: Question text without multiple choice options
+            question_with_options: Full question with options and proper template
+            
+        Returns:
+            Dictionary with context and answer
+        """
+        # Stage 1: Get background context (uses parent class method)
+        context = self.get_background_context(question_without_options)
+        
+        # Stage 2: Answer question with context using Rosetta
+        answer = self.answer_with_context(
+            question=question_with_options,
+            context=context,
+            original_question=question_without_options
+        )
+        
+        return {
+            "context": context,
+            "answer": answer
+        }
+
+
+
+class MultiModalInference:
+    """Multi-modal VLM+LLM inference pipeline."""
+    
+    def __init__(
+        self,
+        vlm_model_path: str,
+        llm_model_path: str,
+        device: str = "cuda",
+        max_new_tokens: int = 1024,
+        generation_config: Optional[Dict[str, Any]] = None
+    ):
+        """
+        Initialize multi-modal pipeline.
+        
+        Args:
+            vlm_model_path: Path to VLM model
+            llm_model_path: Path to LLM model  
+            device: Device to use
+            max_new_tokens: Maximum number of new tokens to generate
+            generation_config: Optional generation configuration to apply to models
+        """
+        self.device = device
+        self.max_new_tokens = max_new_tokens
+        self.generation_config = generation_config or {}
+        self._load_models(vlm_model_path, llm_model_path)
+    
+    def _load_models(self, vlm_path: str, llm_path: str):
+        """Load VLM and LLM models."""
+        # Load VLM
+        self.vlm_model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
+            vlm_path,
+            torch_dtype=torch.bfloat16,
+            device_map={"": self.device},
+        )
+        # Apply generation config to VLM model
+        apply_generation_config(self.vlm_model, self.generation_config)
+        self.vlm_processor = AutoProcessor.from_pretrained(vlm_path)
+        
+        # Load LLM
+        self.llm_tokenizer = AutoTokenizer.from_pretrained(llm_path)
+        self.llm_model = AutoModelForCausalLM.from_pretrained(
+            llm_path, torch_dtype=torch.bfloat16, device_map={"": self.device}
+        )
+        # Apply generation config to LLM model
+        apply_generation_config(self.llm_model, self.generation_config)
+    
+    def get_image_description(
+        self,
+        image_path: str,
+        prompt: str = "Describe this image in detail.",
+        max_new_tokens: Optional[int] = None
+    ) -> str:
+        """
+        Get image description from VLM.
+        
+        Args:
+            image_path: Path to image
+            prompt: Description prompt
+            max_new_tokens: Max tokens to generate (uses instance default if None)
+            
+        Returns:
+            Image description
+        """
+        messages = [{
+            "role": "user",
+            "content": [
+                {"type": "image", "image": image_path},
+                {"type": "text", "text": prompt}
+            ]
+        }]
+        
+        text = self.vlm_processor.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
+        image_inputs, video_inputs = process_vision_info(messages)
+        inputs = self.vlm_processor(
+            text=[text],
+            images=image_inputs,
+            videos=video_inputs,
+            padding=True,
+            return_tensors="pt",
+        )
+        inputs = inputs.to(self.device)
+        
+        if max_new_tokens is None:
+            max_new_tokens = self.max_new_tokens
+            
+        with torch.inference_mode():
+            outputs = self.vlm_model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False,
+            )
+        
+        generated_ids = outputs[:, inputs["input_ids"].shape[-1]:]
+        description = self.vlm_processor.batch_decode(
+            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        
+        return description
+    
+    def answer_with_context(
+        self,
+        question: str,
+        context: str,
+        max_new_tokens: Optional[int] = None,
+        original_question: Optional[str] = None
+    ) -> str:
+        """
+        Answer question using LLM with context.
+        
+        Args:
+            question: Question to answer
+            context: Context (e.g., image description from VLM)
+            max_new_tokens: Max tokens to generate (uses instance default if None)
+            original_question: Original question asked to VLM (for conversation format)
+            
+        Returns:
+            Generated answer
+        """
+        # Use conversation format: user asks about image, assistant describes, user asks follow-up
+        if original_question:
+            messages = [
+                {"role": "user", "content": original_question},
+                {"role": "assistant", "content": context},
+                {"role": "user", "content": question}
+            ]
+        else:
+            # Fallback to simple format
+            messages = [{"role": "user", "content": f"Context: {context}\n\nQuestion: {question}"}]
+
+        template_kwargs = {'enable_thinking': False}
+        
+        # Some tokenizers may not support enable_thinking parameter
+        inputs = self.llm_tokenizer.apply_chat_template(
+            messages,
+            tokenize=True,
+            add_generation_prompt=True,
+            return_tensors="pt",
+            **template_kwargs
+        )
+
+        inputs = inputs.to(self.llm_model.device)
+        
+        if max_new_tokens is None:
+            max_new_tokens = self.max_new_tokens
+        
+        with torch.inference_mode():
+            outputs = self.llm_model.generate(
+                inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False
+            )
+        
+        generated_ids = outputs[:, inputs.shape[-1]:]
+        answer = self.llm_tokenizer.batch_decode(
+            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False
+        )[0]
+        
+        return answer
+    
+    def process(
+        self,
+        image_path: str,
+        question: str,
+        description_prompt: str = "Briefly describe this image."
+    ) -> Dict[str, str]:
+        """
+        Full multi-stage processing.
+        
+        Args:
+            image_path: Path to image
+            question: Question to answer
+            description_prompt: Prompt for image description
+            
+        Returns:
+            Dictionary with description and answer
+        """
+        # Stage 1: Get image description
+        description = self.get_image_description(image_path, description_prompt)
+        
+        # Stage 2: Answer question with context (pass original prompt for conversation format)
+        answer = self.answer_with_context(
+            question=question, 
+            context=description,
+            original_question=description_prompt
+        )
+        
+        return {
+            "description": description,
+            "answer": answer
+        }

rosetta/model/ablation_projector.py

@@ -0,0 +1,291 @@
+"""
+Ablation Projector: A configurable projector for ablation studies based on C2CProjector.
+Allows gradual removal of components to study their individual contributions.
+"""
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Optional, Tuple, Literal
+
+from rosetta.utils.registry import register_model, capture_init_args
+from rosetta.model.projector import Projector
+from rosetta.model.projector import RegularMLP
+
+
+@register_model
+@capture_init_args
+class AblationProjector(Projector):
+    """
+    Ablation study projector based on C2CProjector with configurable component removal.
+    
+    Ablation levels:
+    0. Full C2C (baseline)
+    1. Remove scalar weights (set to 1.0)
+    2. Remove gates (set to 1.0) 
+    3. Remove target contribution (only use source)
+    4. Remove gates only (gates=1.0), keep scalars and target
+    
+    Each level builds on the previous one, allowing gradual degradation study.
+    """
+
+    def __init__(
+        self,
+        source_dim: int,
+        target_dim: int,
+        source_num_heads: int = 1,
+        target_num_heads: int = 1,
+        intermediate_dim: int = 1024,
+        hidden_dim: int = 1024,
+        num_layers: int = 3,
+        dropout: float = 0.1,
+        initial_temperature: float = 1.0,
+        final_temperature: float = 0.001,
+        anneal_steps: int = 1929,
+        dtype: torch.dtype = torch.float32,
+        
+        # Ablation configuration
+        ablation_level: int = 0,  # 0=full, 1=no_scalar, 2=no_gate+no_scalar, 3=no_target, 4=no_gate_only
+        use_scalar_weights: bool = True,  # Can be overridden by ablation_level
+        use_gates: bool = True,          # Can be overridden by ablation_level  
+        use_target: bool = True,         # Can be overridden by ablation_level
+    ):
+        super().__init__()
+
+        assert 0 <= ablation_level <= 4, "ablation_level must be 0, 1, 2, 3, or 4"
+
+        # Dimensions
+        self.source_dim = source_dim
+        self.target_dim = target_dim
+        self.source_num_heads = source_num_heads
+        self.target_num_heads = target_num_heads
+        self.ablation_level = ablation_level
+
+        # Override component usage based on ablation level
+        if ablation_level == 4:
+            # Special case: disable gates only, keep scalars and target
+            use_scalar_weights = True
+            use_gates = False
+            use_target = True
+        else:
+            if ablation_level >= 1:
+                use_scalar_weights = False
+            if ablation_level >= 2: 
+                use_gates = False
+            if ablation_level >= 3:
+                use_target = False
+            
+        self.use_scalar_weights = use_scalar_weights
+        self.use_gates = use_gates
+        self.use_target = use_target
+
+        # Sizes
+        in_dim = source_dim * source_num_heads
+        out_dim = target_dim * target_num_heads
+
+        # 1) concat(source_X, target_X) then project to hidden_dim
+        # If not using target, only use source features
+        if self.use_target:
+            self.key_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype)
+            self.value_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype)
+        else:
+            # Only use source features
+            self.key_in = nn.Linear(in_dim, hidden_dim, bias=True, dtype=dtype)
+            self.value_in = nn.Linear(in_dim, hidden_dim, bias=True, dtype=dtype)
+
+        # 2) one-layer common embedding MLP to get intermediate representation (at hidden_dim)
+        self.key_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype)
+        self.value_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype)
+
+        # 3a) intermediate representation → (L-2)-layer MLP for weights → project to head dim
+        # Only build if using scalar weights
+        if self.use_scalar_weights:
+            self.key_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype)
+            self.value_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype)
+            self.key_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype)
+            self.value_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype)
+
+        # 3b) intermediate representation → (L-2)-layer MLP for projected_X → finally project hidden_dim → out_dim
+        self.key_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype)
+        self.value_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype)
+        self.key_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype)
+        self.value_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype)
+
+        # Scalar key/value gate parameters and temperature schedule
+        # Only build if using gates
+        if self.use_gates:
+            self.key_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype))
+            self.value_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype))
+            self.use_gumbel = True
+            self.register_buffer("gate_temperature", torch.tensor(initial_temperature, dtype=dtype))
+            self.initial_temperature = initial_temperature
+            self.final_temperature = final_temperature
+            self.anneal_steps = anneal_steps
+
+        # Temperature for weight normalization
+        self.scalar_temperature = 1.0
+
+    def update_temperature(self, step: int):
+        """Update temperature using exponential annealing schedule for gates."""
+        if self.use_gates:
+            ratio = min(step / self.anneal_steps, 1.0)
+            temp = self.initial_temperature * (self.final_temperature / self.initial_temperature) ** ratio
+            self.gate_temperature.fill_(temp)
+
+    def forward(
+        self,
+        source_kv: Tuple[Tensor, Tensor],
+        target_kv: Tuple[Tensor, Tensor],
+        position_ids: Optional[Tensor] = None,
+        max_pos: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        source_key, source_value = source_kv
+        target_key, target_value = target_kv
+
+        B, Hs, N, Ds = source_key.shape
+        _, Ht, _, Dt = target_key.shape
+
+        # Flatten heads
+        source_key_flat = source_key.transpose(1, 2).contiguous().view(B, N, Hs * Ds)
+        source_value_flat = source_value.transpose(1, 2).contiguous().view(B, N, Hs * Ds)
+        target_key_flat = target_key.transpose(1, 2).contiguous().view(B, N, Ht * Dt)
+        target_value_flat = target_value.transpose(1, 2).contiguous().view(B, N, Ht * Dt)
+
+        # 1) Prepare input features based on ablation level
+        if self.use_target:
+            # Full C2C: concat source and target features
+            key_cat = torch.cat([source_key_flat, target_key_flat], dim=-1)
+            value_cat = torch.cat([source_value_flat, target_value_flat], dim=-1)
+        else:
+            # Ablation level 3: only use source features
+            key_cat = source_key_flat
+            value_cat = source_value_flat
+
+        # 2) project to hidden dim
+        key_hidden = self.key_in(key_cat)
+        value_hidden = self.value_in(value_cat)
+
+        # 3) one-layer common embedding MLP to get intermediate representation (at hidden_dim)
+        key_hidden = self.key_mlp1(key_hidden)
+        value_hidden = self.value_mlp1(value_hidden)
+
+        # 4b) intermediate representation -> projected feature path
+        key_proj_hidden = self.key_proj_out(self.key_proj_mlp2(key_hidden)) # (B, N, Ht * Dt)
+        value_proj_hidden = self.value_proj_out(self.value_proj_mlp2(value_hidden)) # (B, N, Ht * Dt)
+        projected_key = key_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt)
+        projected_value = value_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt)
+
+        # 4a) intermediate representation -> scalar path (if using scalar weights)
+        if self.use_scalar_weights:
+            key_scalar = self.key_scalar_head(self.key_scalar_mlp2(key_hidden))       # (B, N, Ht)
+            value_scalar = self.value_scalar_head(self.value_scalar_mlp2(value_hidden)) # (B, N, Ht)
+            key_scalar = key_scalar.permute(0, 2, 1).unsqueeze(-1)   # (B, Ht, N, 1)
+            value_scalar = value_scalar.permute(0, 2, 1).unsqueeze(-1)  # (B, Ht, N, 1)
+            # Normalize scalars
+            norm_key_scalar = torch.sigmoid(key_scalar)
+            norm_value_scalar = torch.sigmoid(value_scalar)
+        else:
+            # Ablation level 1+: set scalar weights to 1.0
+            norm_key_scalar = torch.ones(B, Ht, N, 1, device=projected_key.device, dtype=projected_key.dtype)
+            norm_value_scalar = torch.ones(B, Ht, N, 1, device=projected_value.device, dtype=projected_value.dtype)
+
+        # Key/value gates (if using gates)
+        if self.use_gates:
+            key_gate_logit = self.key_gate_logit.view(1, 1, 1, 1)
+            value_gate_logit = self.value_gate_logit.view(1, 1, 1, 1)
+            if self.training and self.use_gumbel:
+                u1 = torch.rand(B, Ht, N, 1, device=key_gate_logit.device, dtype=key_gate_logit.dtype)
+                u2 = torch.rand(B, Ht, N, 1, device=value_gate_logit.device, dtype=value_gate_logit.dtype)
+                g1 = -torch.log(-torch.log(u1 + 1e-20) + 1e-20)
+                g2 = -torch.log(-torch.log(u2 + 1e-20) + 1e-20)
+                key_gate = torch.sigmoid((key_gate_logit + g1) / self.gate_temperature)
+                value_gate = torch.sigmoid((value_gate_logit + g2) / self.gate_temperature)
+            else:
+                key_gate = (key_gate_logit > 0).float()
+                value_gate = (value_gate_logit > 0).float()
+        else:
+            # Gates disabled: set gates to 1.0 (always open)
+            key_gate = torch.ones(B, Ht, N, 1, device=projected_key.device, dtype=projected_key.dtype)
+            value_gate = torch.ones(B, Ht, N, 1, device=projected_value.device, dtype=projected_value.dtype)
+
+        # Compute projected contribution
+        projected_key_term = key_gate * norm_key_scalar * projected_key
+        projected_value_term = value_gate * norm_value_scalar * projected_value
+
+        # Compute target contribution (if using target)
+        if self.use_target:
+            # Full C2C: add target with projected
+            output_key = target_key + projected_key_term
+            output_value = target_value + projected_value_term
+        else:
+            # Ablation level 3: only use projected (no target)
+            output_key = projected_key_term
+            output_value = projected_value_term
+
+        return output_key, output_value
+
+    def get_ablation_info(self) -> dict:
+        """Return information about current ablation configuration."""
+        return {
+            'ablation_level': self.ablation_level,
+            'use_scalar_weights': self.use_scalar_weights,
+            'use_gates': self.use_gates,
+            'use_target': self.use_target,
+            'description': self._get_ablation_description()
+        }
+    
+    def _get_ablation_description(self) -> str:
+        """Get human-readable description of current ablation level."""
+        descriptions = {
+            0: "Full C2C (baseline)",
+            1: "No scalar weights (scalars=1.0)",
+            2: "No gates (gates=1.0) + No scalar weights",
+            3: "No target (source-only) + No gates + No scalar weights",
+            4: "No gates (gates=1.0), keep scalars and target"
+        }
+        return descriptions.get(self.ablation_level, "Unknown ablation level")
+
+
+# Convenience functions for creating specific ablation levels
+def create_ablation_projector(
+    source_dim: int,
+    target_dim: int,
+    source_num_heads: int = 1,
+    target_num_heads: int = 1,
+    ablation_level: int = 0,
+    **kwargs
+) -> AblationProjector:
+    """Create an AblationProjector with specified ablation level."""
+    return AblationProjector(
+        source_dim=source_dim,
+        target_dim=target_dim,
+        source_num_heads=source_num_heads,
+        target_num_heads=target_num_heads,
+        ablation_level=ablation_level,
+        **kwargs
+    )
+
+
+def create_full_c2c_projector(**kwargs) -> AblationProjector:
+    """Create full C2C projector (ablation level 0)."""
+    return create_ablation_projector(ablation_level=0, **kwargs)
+
+
+def create_no_scalar_projector(**kwargs) -> AblationProjector:
+    """Create projector without scalar weights (ablation level 1)."""
+    return create_ablation_projector(ablation_level=1, **kwargs)
+
+
+def create_no_gate_projector(**kwargs) -> AblationProjector:
+    """Create projector without gates (ablation level 2)."""
+    return create_ablation_projector(ablation_level=2, **kwargs)
+
+
+def create_source_only_projector(**kwargs) -> AblationProjector:
+    """Create source-only projector (ablation level 3)."""
+    return create_ablation_projector(ablation_level=3, **kwargs)
+
+
+def create_no_gate_only_projector(**kwargs) -> AblationProjector:
+    """Create projector without gates but with scalar weights and target (ablation level 4)."""
+    return create_ablation_projector(ablation_level=4, **kwargs)

rosetta/model/aligner.py

@@ -0,0 +1,639 @@
+"""
+Token Aligner for handling different tokenizers between SLM and LLM models.
+
+This module provides functionality to align tokens between two different tokenizers,
+handling cases where the same text is tokenized differently.
+"""
+
+from typing import List, Tuple, Optional, Dict, Literal, Union
+import torch
+from transformers import PreTrainedTokenizerBase
+from enum import Enum
+
+
+class AlignmentStrategy(Enum):
+    """Strategies for handling 1-to-many token alignments"""
+    FIRST = "first"    # Always take the first LLM token
+    LONGEST = "longest"  # Take the LLM token with the longest string
+
+
+class TokenAligner:
+    """
+    Aligns tokens between SLM (Small Language Model) and LLM (Large Language Model) tokenizers.
+    
+    This class handles the case where the same text sequence is tokenized differently
+    by different tokenizers, using the SLM tokenization as the base and finding
+    corresponding LLM tokens for each SLM token.
+    """
+    
+    def __init__(
+        self,
+        slm_tokenizer: PreTrainedTokenizerBase,
+        llm_tokenizer: PreTrainedTokenizerBase,
+        strategy: Union[AlignmentStrategy, str] = AlignmentStrategy.FIRST,
+        verbose: bool = False
+    ):
+        """
+        Initialize the TokenAligner.
+        
+        Args:
+            slm_tokenizer: The tokenizer for the Small Language Model (base)
+            llm_tokenizer: The tokenizer for the Large Language Model
+            strategy: Strategy for handling 1-to-many token mappings
+                     Either AlignmentStrategy enum or string ('first' or 'longest')
+            verbose: Whether to print debug information during alignment
+        """
+        self.slm_tokenizer = slm_tokenizer
+        self.llm_tokenizer = llm_tokenizer
+        
+        if self.slm_tokenizer.pad_token is None:
+            self.slm_tokenizer.pad_token = self.slm_tokenizer.eos_token
+            self.slm_tokenizer.pad_token_id = self.slm_tokenizer.eos_token_id
+        if self.llm_tokenizer.pad_token is None:
+            self.llm_tokenizer.pad_token = self.llm_tokenizer.eos_token
+            self.llm_tokenizer.pad_token_id = self.llm_tokenizer.eos_token_id
+
+        # Handle string strategy input
+        if isinstance(strategy, str):
+            strategy = AlignmentStrategy(strategy.lower())
+        self.strategy = strategy
+        self.verbose = verbose
+        
+        # Cache for token mappings to improve performance
+        self._alignment_cache: Dict[Tuple[int, ...], List[int]] = {}
+    
+    def align_tokens(
+        self,
+        slm_token_ids: Union[List[int], torch.Tensor],
+        return_mapping: bool = False
+    ) -> Union[List[int], Tuple[List[int], List[Tuple[int, List[int]]]]]:
+        """
+        Align SLM tokens to LLM tokens.
+        
+        Args:
+            slm_token_ids: Token IDs from the SLM tokenizer
+            return_mapping: If True, also return the detailed mapping
+        
+        Returns:
+            If return_mapping is False: List of aligned LLM token IDs
+            If return_mapping is True: Tuple of (aligned_llm_token_ids, mapping_details)
+                where mapping_details is a list of (slm_token_id, [candidate_llm_token_ids])
+        """
+        # Convert to list if tensor
+        if isinstance(slm_token_ids, torch.Tensor):
+            slm_token_ids = slm_token_ids.tolist()
+        
+        # Check cache
+        cache_key = tuple(slm_token_ids)
+        if cache_key in self._alignment_cache and not return_mapping:
+            return self._alignment_cache[cache_key]
+        
+        aligned_llm_tokens = []
+        mapping_details = []
+        
+        for slm_token_id in slm_token_ids:
+            # Decode SLM token to string (without special token processing)
+            slm_token_str = self.slm_tokenizer.decode(
+                [slm_token_id], 
+                skip_special_tokens=False,
+                clean_up_tokenization_spaces=False
+            )
+            
+            # Handle special tokens
+            if slm_token_id in self.slm_tokenizer.all_special_ids:
+                # Try to find corresponding special token in LLM tokenizer
+                llm_token_id = self._map_special_token(slm_token_id, slm_token_str)
+                aligned_llm_tokens.append(llm_token_id)
+                mapping_details.append((slm_token_id, [llm_token_id]))
+                continue
+            
+            # Tokenize the string with LLM tokenizer
+            llm_token_ids = self.llm_tokenizer.encode(
+                slm_token_str,
+                add_special_tokens=False,
+                return_tensors=None
+            )
+            
+            if len(llm_token_ids) == 0:
+                # Handle empty tokenization (shouldn't normally happen)
+                if self.verbose:
+                    print(f"Warning: SLM token {slm_token_id} ('{slm_token_str}') "
+                          f"resulted in empty LLM tokenization")
+                # Use unknown token as fallback
+                llm_token_id = self.llm_tokenizer.unk_token_id or 0
+                aligned_llm_tokens.append(llm_token_id)
+                mapping_details.append((slm_token_id, [llm_token_id]))
+                
+            elif len(llm_token_ids) == 1:
+                # Perfect 1-to-1 mapping
+                aligned_llm_tokens.append(llm_token_ids[0])
+                mapping_details.append((slm_token_id, llm_token_ids))
+                
+            else:
+                # 1-to-many mapping, apply strategy
+                selected_token = self._apply_strategy(llm_token_ids, slm_token_str)
+                aligned_llm_tokens.append(selected_token)
+                mapping_details.append((slm_token_id, llm_token_ids))
+                
+                if self.verbose:
+                    selected_str = self.llm_tokenizer.decode(
+                        [selected_token],
+                        skip_special_tokens=False,
+                        clean_up_tokenization_spaces=False
+                    )
+                    print(f"SLM token {slm_token_id} ('{slm_token_str}') -> "
+                          f"LLM tokens {llm_token_ids}, selected {selected_token} ('{selected_str}')")
+        
+        # Cache the result
+        self._alignment_cache[cache_key] = aligned_llm_tokens
+        
+        if return_mapping:
+            return aligned_llm_tokens, mapping_details
+        return aligned_llm_tokens
+    
+    def _map_special_token(self, slm_token_id: int, slm_token_str: str) -> int:
+        """
+        Map special tokens between tokenizers.
+        
+        Args:
+            slm_token_id: The SLM special token ID
+            slm_token_str: The string representation of the special token
+        
+        Returns:
+            The corresponding LLM token ID
+        """
+        # Common special token mappings
+        special_token_map = {
+            self.slm_tokenizer.pad_token_id: self.llm_tokenizer.pad_token_id,
+            self.slm_tokenizer.eos_token_id: self.llm_tokenizer.eos_token_id,
+            self.slm_tokenizer.bos_token_id: self.llm_tokenizer.bos_token_id,
+            self.slm_tokenizer.unk_token_id: self.llm_tokenizer.unk_token_id,
+        }
+        
+        # Direct mapping if available
+        if slm_token_id in special_token_map and special_token_map[slm_token_id] is not None:
+            return special_token_map[slm_token_id]
+        
+        # Try to find by string representation
+        try:
+            llm_token_id = self.llm_tokenizer.convert_tokens_to_ids(slm_token_str)
+            if llm_token_id != self.llm_tokenizer.unk_token_id:
+                return llm_token_id
+        except:
+            pass
+        
+        # Fallback to unknown token
+        return self.llm_tokenizer.unk_token_id or 0
+    
+    def _apply_strategy(self, llm_token_ids: List[int], original_str: str) -> int:
+        """
+        Apply the selected strategy to choose one LLM token from multiple candidates.
+        
+        Args:
+            llm_token_ids: List of candidate LLM token IDs
+            original_str: The original string from SLM token
+        
+        Returns:
+            The selected LLM token ID
+        """
+        if self.strategy == AlignmentStrategy.FIRST:
+            return llm_token_ids[0]
+        
+        elif self.strategy == AlignmentStrategy.LONGEST:
+            # Find the token with the longest string representation
+            longest_token = llm_token_ids[0]
+            longest_length = 0
+            
+            for token_id in llm_token_ids:
+                token_str = self.llm_tokenizer.decode(
+                    [token_id],
+                    skip_special_tokens=False,
+                    clean_up_tokenization_spaces=False
+                )
+                if len(token_str) > longest_length:
+                    longest_length = len(token_str)
+                    longest_token = token_id
+            
+            return longest_token
+        
+        else:
+            # Default to first token if unknown strategy
+            return llm_token_ids[0]
+    
+    def align_sequence(
+        self,
+        text: str,
+        return_details: bool = False
+    ) -> Union[Tuple[List[int], List[int]], Dict[str, any]]:
+        """
+        Tokenize text with both tokenizers and return aligned sequences.
+        
+        Args:
+            text: The input text to tokenize and align
+            return_details: If True, return detailed alignment information
+        
+        Returns:
+            If return_details is False: Tuple of (slm_token_ids, aligned_llm_token_ids)
+            If return_details is True: Dictionary with detailed alignment information
+        """
+        # Tokenize with SLM
+        slm_tokens = self.slm_tokenizer.encode(
+            text,
+            add_special_tokens=True,
+            return_tensors=None
+        )
+        
+        # Get aligned LLM tokens
+        if return_details:
+            aligned_llm_tokens, mapping = self.align_tokens(slm_tokens, return_mapping=True)
+            
+            # Decode tokens for inspection
+            slm_decoded = [
+                self.slm_tokenizer.decode([tid], skip_special_tokens=False, clean_up_tokenization_spaces=False)
+                for tid in slm_tokens
+            ]
+            llm_decoded = [
+                self.llm_tokenizer.decode([tid], skip_special_tokens=False, clean_up_tokenization_spaces=False)
+                for tid in aligned_llm_tokens
+            ]
+            
+            # Original LLM tokenization for comparison
+            original_llm_tokens = self.llm_tokenizer.encode(
+                text,
+                add_special_tokens=True,
+                return_tensors=None
+            )
+            
+            # One-to-one mapping statistics
+            num_tokens = len(slm_tokens)
+            one_to_one_count = sum(1 for _slm_id, candidates in mapping if len(candidates) == 1)
+            one_to_one_rate = (one_to_one_count / num_tokens) if num_tokens > 0 else 0.0
+            
+            return {
+                'text': text,
+                'slm_token_ids': slm_tokens,
+                'slm_decoded': slm_decoded,
+                'aligned_llm_token_ids': aligned_llm_tokens,
+                'aligned_llm_decoded': llm_decoded,
+                'original_llm_token_ids': original_llm_tokens,
+                'mapping': mapping,
+                'strategy': self.strategy.value,
+                'num_tokens': num_tokens,
+                'one_to_one_count': one_to_one_count,
+                'one_to_one_rate': one_to_one_rate
+            }
+        else:
+            aligned_llm_tokens = self.align_tokens(slm_tokens)
+            return slm_tokens, aligned_llm_tokens
+    
+    def visualize_alignment(self, text: str):
+        """
+        Print a visual representation of the token alignment.
+        
+        Args:
+            text: The text to analyze
+        """
+        details = self.align_sequence(text, return_details=True)
+        
+        print("=" * 80)
+        print(f"Text: {text}")
+        print(f"Strategy: {details['strategy']}")
+        print("=" * 80)
+        print(f"SLM tokens ({len(details['slm_token_ids'])}): {details['slm_token_ids']}")
+        print(f"Aligned LLM tokens ({len(details['aligned_llm_token_ids'])}): {details['aligned_llm_token_ids']}")
+        print(f"Original LLM tokens ({len(details['original_llm_token_ids'])}): {details['original_llm_token_ids']}")
+        print("-" * 80)
+        print("Token-by-token alignment:")
+        
+        for i, (slm_id, llm_id) in enumerate(zip(details['slm_token_ids'], details['aligned_llm_token_ids'])):
+            slm_str = details['slm_decoded'][i]
+            llm_str = details['aligned_llm_decoded'][i]
+            mapping_info = details['mapping'][i]
+            
+            if len(mapping_info[1]) > 1:
+                candidates_str = ', '.join([
+                    f"{tid}:'{self.llm_tokenizer.decode([tid], skip_special_tokens=False, clean_up_tokenization_spaces=False)}'"
+                    for tid in mapping_info[1]
+                ])
+                print(f"  [{i:3d}] SLM {slm_id:6d} ('{slm_str}') -> "
+                      f"LLM {llm_id:6d} ('{llm_str}') "
+                      f"[candidates: {candidates_str}]")
+            else:
+                print(f"  [{i:3d}] SLM {slm_id:6d} ('{slm_str}') -> "
+                      f"LLM {llm_id:6d} ('{llm_str}')")
+        print("=" * 80)
+    
+    def clear_cache(self):
+        """Clear the alignment cache."""
+        self._alignment_cache.clear()
+
+    # ========================
+    # Chat messages alignment
+    # ========================
+    def _apply_chat_template_to_ids(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        messages: List[Dict[str, str]],
+        add_generation_prompt: bool,
+        enable_thinking: bool,
+        remove_last_surfix: bool
+    ) -> Tuple[str, List[int], Optional[List[Tuple[int, int]]]]:
+        """
+        Apply chat template (no tokenization) then tokenize to ids with optional offsets.
+        If remove_last_surfix is True, remove the last suffix from the LLM text
+        Returns (templated_text, input_ids, offsets) where offsets may be None.
+        """
+        if remove_last_surfix:
+            assert messages[-1]["role"] == "assistant", "Last message must be an assistant message"
+            templated_text = tokenizer.apply_chat_template(
+                messages[:-1],
+                tokenize=False,
+                add_generation_prompt=True,
+                enable_thinking=enable_thinking
+            )
+            templated_text += messages[-1]["content"]
+        else:
+            templated_text = tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=add_generation_prompt,
+                enable_thinking=enable_thinking
+            )
+        encoded = tokenizer(
+            templated_text,
+            add_special_tokens=False,
+            return_offsets_mapping=True
+        )
+        input_ids: List[int] = encoded["input_ids"]
+        offsets = encoded.get("offset_mapping")
+        return templated_text, input_ids, offsets
+
+    @staticmethod
+    def _first_non_empty_content(messages: List[Dict[str, str]]) -> Optional[str]:
+        for m in messages:
+            content = m.get("content")
+            if isinstance(content, str) and len(content.strip()) > 0:
+                return content
+        return None
+
+    def _find_boundary_token_index(
+        self,
+        tokenizer: PreTrainedTokenizerBase,
+        templated_text: str,
+        offsets: Optional[List[Tuple[int, int]]],
+        content_text: Optional[str]
+    ) -> int:
+        """
+        Find token index where the first non-empty message content starts.
+        Falls back to 0 if not found.
+        """
+        if not content_text:
+            return 0
+        char_idx = templated_text.find(content_text)
+        if char_idx < 0:
+            # Try a shorter probe to improve chances
+            probe = content_text[: min(32, len(content_text))]
+            if len(probe) > 0:
+                char_idx = templated_text.find(probe)
+        if char_idx < 0:
+            return 0
+
+        if offsets:
+            for idx, (start, _end) in enumerate(offsets):
+                if start >= char_idx:
+                    return idx
+            return len(offsets)
+
+        # Fallback without offsets: tokenize prefix and count tokens
+        prefix = templated_text[:char_idx]
+        prefix_ids = tokenizer(prefix, add_special_tokens=False)["input_ids"]
+        return len(prefix_ids)
+
+    @staticmethod
+    def _compute_content_spans(templated_text: str, messages: List[Dict[str, str]]) -> List[Tuple[int, int]]:
+        """
+        Compute character spans in templated_text that correspond to message contents.
+        Searches sequentially to reduce ambiguity when contents repeat.
+        Enhanced matching: ensures the found content is followed by '<' (special token start)
+        to avoid matching content inside special tokens like <begin_of_text>.
+        """
+        spans: List[Tuple[int, int]] = []
+        search_from = 0
+        for m in messages:
+            content = m.get("content")
+            if not isinstance(content, str) or len(content) == 0:
+                continue
+            
+            # Find all possible matches starting from search_from
+            idx = search_from
+            found_valid_match = False
+            
+            while idx < len(templated_text):
+                idx = templated_text.find(content, idx)
+                if idx < 0:
+                    break
+                
+                # Check if this match is valid (followed by '<' indicating a special token)
+                end_pos = idx + len(content)
+                if end_pos < len(templated_text) and templated_text[end_pos] == '<':
+                    # Valid match: content is followed by a special token
+                    spans.append((idx, end_pos))
+                    search_from = end_pos
+                    found_valid_match = True
+                    break
+                else:
+                    # Check if this is the end of the text (also valid for last message)
+                    if end_pos == len(templated_text):
+                        spans.append((idx, end_pos))
+                        search_from = end_pos
+                        found_valid_match = True
+                        break
+                
+                # Invalid match, try next occurrence
+                idx += 1
+            
+            # Fallback: if no valid match found with '<' requirement, use the old method
+            # but only as a last resort and with additional validation
+            if not found_valid_match:
+                idx = templated_text.find(content, search_from)
+                if idx < 0:
+                    # Try searching from start as last resort
+                    idx = templated_text.find(content)
+                
+                if idx >= 0:
+                    end_pos = idx + len(content)
+                    # Additional check: avoid matching inside obvious special tokens
+                    # Check if we're inside a special token (preceded by '<' and not followed by '>')
+                    start_context = templated_text[max(0, idx-10):idx]
+                    end_context = templated_text[end_pos:min(len(templated_text), end_pos+10)]
+                    
+                    # Skip if we're clearly inside a special token
+                    if ('<' in start_context and '>' not in start_context and 
+                        'begin_of_text' in templated_text[max(0, idx-20):idx+20]):
+                        # This looks like we're matching inside <begin_of_text> or similar
+                        continue
+                    
+                    spans.append((idx, end_pos))
+                    search_from = end_pos
+        
+        return spans
+
+    @staticmethod
+    def _build_token_mask_from_spans(
+        offsets: Optional[List[Tuple[int, int]]],
+        num_tokens: int,
+        spans: List[Tuple[int, int]]
+    ) -> List[bool]:
+        """
+        Build a boolean mask for tokens whose offset range overlaps any span.
+        If offsets are missing, default to all False.
+        """
+        if not offsets or len(offsets) != num_tokens:
+            return [False] * num_tokens
+        mask: List[bool] = []
+        for (start, end) in offsets:
+            if end <= start:
+                mask.append(False)
+                continue
+            is_msg = False
+            for s, e in spans:
+                # overlap check
+                if start < e and end > s:
+                    is_msg = True
+                    break
+            mask.append(is_msg)
+        return mask
+
+    @staticmethod
+    def _spans_to_token_ranges(
+        offsets: List[Tuple[int, int]],
+        spans: List[Tuple[int, int]]
+    ) -> List[Tuple[int, int]]:
+        """
+        Convert character spans to token index ranges using offsets.
+        start token = first token with end > span_start
+        end token = first token with start >= span_end
+        """
+        ranges: List[Tuple[int, int]] = []
+        n = len(offsets)
+        for s, e in spans:
+            # find start index
+            start_idx = 0
+            while start_idx < n and offsets[start_idx][1] <= s:
+                start_idx += 1
+            # find end index
+            end_idx = start_idx
+            while end_idx < n and offsets[end_idx][0] < e:
+                end_idx += 1
+            ranges.append((start_idx, end_idx))
+        return ranges
+
+    def align_chat_messages(
+        self,
+        messages: List[Dict[str, str]],
+        add_generation_prompt: bool = True,
+        enable_thinking: bool = False,
+        return_details: bool = False,
+        remove_last_surfix: bool = False
+    ) -> Dict[str, any]:
+        """
+        Align chat-templated sequences by sections (template/message/template...):
+        - Preserve all template tokens (pad the shorter template section)
+        - For each message section, map SLM tokens to LLM tokens 1:1 via strategy
+        - If remove_last_surfix is True, remove the last suffix from the LLM text
+        Returns essentials: slm_ids_padded, llm_ids_padded, message_mask (shared),
+        slm_padding_mask, llm_padding_mask (True where token is padding inserted).
+        When return_details=True, also returns 'sections' with aligned ranges.
+        """
+        assert not (add_generation_prompt and remove_last_surfix), "add_generation_prompt and remove_last_surfix cannot be True at the same time"
+
+        # Build templated sequences with offsets
+        slm_text, slm_ids, slm_offsets = self._apply_chat_template_to_ids(
+            self.slm_tokenizer, messages, add_generation_prompt, enable_thinking, remove_last_surfix
+        )
+        llm_text, llm_ids, llm_offsets = self._apply_chat_template_to_ids(
+            self.llm_tokenizer, messages, add_generation_prompt, enable_thinking, remove_last_surfix
+        )
+
+        # Required pad tokens
+        assert self.slm_tokenizer.pad_token_id is not None, "SLM pad_token_id required"
+        assert self.llm_tokenizer.pad_token_id is not None, "LLM pad_token_id required"
+        slm_pad_id = self.slm_tokenizer.pad_token_id
+        llm_pad_id = self.llm_tokenizer.pad_token_id
+
+        # Content spans (char) and token ranges
+        content_spans_slm = self._compute_content_spans(slm_text, messages)
+        content_spans_llm = self._compute_content_spans(llm_text, messages)
+        assert slm_offsets is not None and llm_offsets is not None, "offset_mapping required"
+        slm_msg_ranges = self._spans_to_token_ranges(slm_offsets, content_spans_slm)
+        llm_msg_ranges = self._spans_to_token_ranges(llm_offsets, content_spans_llm)
+        # Build section ranges (template/message alternating)
+        def build_sections(total_len: int, msg_ranges: List[Tuple[int,int]]):
+            sections: List[Tuple[str,int,int]] = []
+            prev = 0
+            for (s, e) in msg_ranges:
+                if prev < s:
+                    sections.append(("template", prev, s))
+                sections.append(("message", s, e))
+                prev = e
+            if prev < total_len:
+                sections.append(("template", prev, total_len))
+            return sections
+        slm_sections = build_sections(len(slm_ids), slm_msg_ranges)
+        llm_sections = build_sections(len(llm_ids), llm_msg_ranges)
+        assert len(slm_sections) == len(llm_sections), "Section count mismatch"
+
+        slm_out: List[int] = []
+        llm_out: List[int] = []
+        mask_out: List[bool] = []
+        slm_pad_mask_out: List[bool] = []
+        llm_pad_mask_out: List[bool] = []
+        detailed_sections: List[Dict[str, Union[str, Tuple[int,int]]]] = []
+
+        for (stype_s, s_s, e_s), (stype_l, s_l, e_l) in zip(slm_sections, llm_sections):
+            assert stype_s == stype_l, "Section type mismatch"
+            slm_start_out = len(slm_out)
+            llm_start_out = len(llm_out)
+            if stype_s == "template":
+                slm_seg_len = e_s - s_s
+                llm_seg_len = e_l - s_l
+                target_len = slm_seg_len if slm_seg_len >= llm_seg_len else llm_seg_len
+                slm_pad_needed = target_len - slm_seg_len
+                llm_pad_needed = target_len - llm_seg_len
+                slm_seg = slm_ids[s_s:e_s] + [slm_pad_id] * slm_pad_needed
+                llm_seg = llm_ids[s_l:e_l] + [llm_pad_id] * llm_pad_needed
+                slm_out.extend(slm_seg)
+                llm_out.extend(llm_seg)
+                mask_out.extend([False] * target_len)
+                slm_pad_mask_out.extend([False] * slm_seg_len + [True] * slm_pad_needed)
+                llm_pad_mask_out.extend([False] * llm_seg_len + [True] * llm_pad_needed)
+            else:  # message
+                slm_msg = slm_ids[s_s:e_s]
+                llm_msg = self.align_tokens(slm_msg)
+                assert len(llm_msg) == len(slm_msg)
+                slm_out.extend(slm_msg)
+                llm_out.extend(llm_msg)
+                mask_out.extend([True] * len(slm_msg))
+                # no padding in message sections
+                slm_pad_mask_out.extend([False] * len(slm_msg))
+                llm_pad_mask_out.extend([False] * len(slm_msg))
+            slm_end_out = len(slm_out)
+            llm_end_out = len(llm_out)
+            detailed_sections.append({
+                'type': stype_s,
+                'slm_range': (slm_start_out, slm_end_out),
+                'llm_range': (llm_start_out, llm_end_out)
+            })
+
+        result_min = {
+            'slm_ids_padded': slm_out,
+            'llm_ids_padded': llm_out,
+            'message_mask': mask_out,
+            'slm_padding_mask': slm_pad_mask_out,
+            'llm_padding_mask': llm_pad_mask_out
+        }
+        if return_details:
+            result_min['sections'] = detailed_sections
+            result_min['slm_text'] = slm_text
+            result_min['llm_text'] = llm_text
+        return result_min

rosetta/model/oracle.py

@@ -0,0 +1,562 @@
+from typing import List, Optional, Union
+import torch
+from torch import nn
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+import json
+
+from rosetta.model.projector import Projector
+from rosetta.model.sampling import sample_token
+from transformers.utils import ModelOutput
+try:
+    from transformers.generation.utils import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput
+except Exception:
+    GreedySearchDecoderOnlyOutput = None
+    SampleDecoderOnlyOutput = None
+from rosetta.model.wrapper import RosettaModel
+
+class OracleRosettaModel(nn.Module):
+    """
+    Drop in replacement for the standard transformers LLM models, like Qwen3ForCausalLM
+    """
+    def __init__(self, model_list: List[PreTrainedModel], base_model_idx = 0, projector_list: List[Projector] = [], aggregator_list: List[nn.Module] = []):
+        super().__init__()
+        # model list: a list of model, model 0 by default is the base model
+        # projector list: a list of projector
+        # standard init with additional model list parameter
+        # kv-cache dict: key (source_model_idx, target_model_idx), value (Cache), assume only convert at prefill with one type of model
+        # projector dict: key (source_model_idx, target_model_idx) value dict(key (source_model_layer_idx, M_target value )
+
+        self.base_model_idx = base_model_idx
+        self.model_list = nn.ModuleList(model_list)
+
+        device = model_list[base_model_idx].device
+        dtype = model_list[base_model_idx].dtype
+        self.projector_list = nn.ModuleList(projector_list).to(device=device, dtype=dtype)
+        self.aggregator_list = nn.ModuleList(aggregator_list).to(device=device, dtype=dtype)
+
+        self.projector_dict = {}
+        self.aggregator_dict = {}
+        self.kv_cache_dict = {}
+        self._generation_hook_handlers = []
+
+    @property
+    def device(self):
+        return self.model_list[self.base_model_idx].device
+    
+    def to(self, device):
+        """
+        Move the RosettaModel and all underlying models and projectors to the specified device.
+        """
+        super().to(device)
+        for model in self.model_list:
+            model.to(device)
+        for projector in self.projector_list:
+            projector.to(device)
+        for aggregator in self.aggregator_list:
+            aggregator.to(device)
+        return self
+        
+    # set projector 
+    def set_projector_config(self, 
+                        source_model_idx: int, 
+                        source_model_layer_idx: int, 
+                        target_model_idx: int,
+                        target_model_layer_idx: int, 
+                        projector_idx: int):
+        """
+        Set the projector configuration
+        Args:
+            source_model_idx: int, the index of the source model
+            source_model_layer_idx: int, the index of the source model layer
+            target_model_idx: int, the index of the target model
+            target_model_layer_idx: int, the index of the target model layer
+            projector_idx: int, the index of the projector
+
+        The projector dict structure supports multiple projectors per target layer.
+        Structure:
+        {
+            target_model_idx: {
+                source_model_idx: {
+                    target_model_layer_idx: [(source_model_layer_idx, projector_idx), ...]
+                }
+            }
+        }
+        Repeated calls for the same (target, source, target_layer) append additional pairs.
+        """
+
+        if target_model_idx not in self.projector_dict.keys():
+            self.projector_dict[target_model_idx] = {}
+        if source_model_idx not in self.projector_dict[target_model_idx].keys():
+            self.projector_dict[target_model_idx][source_model_idx] = {}
+        # Accumulate list of (source_layer, projector_idx) for this target layer
+        layer_entry = self.projector_dict[target_model_idx][source_model_idx].get(target_model_layer_idx)
+        if layer_entry is None:
+            self.projector_dict[target_model_idx][source_model_idx][target_model_layer_idx] = [(source_model_layer_idx, projector_idx)]
+        else:
+            layer_entry.append((source_model_layer_idx, projector_idx))
+
+
+    def load_projector(self, projector_list):
+        self.projector_list: List[Projector] = projector_list
+
+    def load_aggregator(self, aggregator_list):
+        self.aggregator_list: List[nn.Module] = aggregator_list
+
+
+    def get_projector(self, 
+                        source_model_idx, 
+                        source_model_layer_idx, 
+                        target_model_idx,
+                        target_model_layer_idx):
+        pair_list = self.projector_dict[target_model_idx][source_model_idx][target_model_layer_idx]
+        if len(pair_list) == 0:
+            raise ValueError("No projector configured for the given target layer")
+        # Prefer exact source layer match
+        for src_layer, projector_id in pair_list:
+            if src_layer == source_model_layer_idx:
+                return self.projector_list[projector_id]
+        # Fallback: return the first projector
+        return self.projector_list[pair_list[0][1]]
+
+    def set_aggregator_idx(self,
+                           source_model_idx: int,
+                           target_model_idx: int,
+                           target_model_layer_idx: int,
+                           aggregator_idx: int):
+        if target_model_idx not in self.aggregator_dict:
+            self.aggregator_dict[target_model_idx] = {}
+        if source_model_idx not in self.aggregator_dict[target_model_idx]:
+            self.aggregator_dict[target_model_idx][source_model_idx] = {}
+        self.aggregator_dict[target_model_idx][source_model_idx][target_model_layer_idx] = aggregator_idx
+
+
+    @staticmethod
+    def load_json(file_name):
+        with open(file_name, "r") as f:
+            result = json.load(f)
+        return result
+    
+    @staticmethod
+    def _convert_dict_keys_to_ints(obj):
+        """
+        Recursively convert dictionary keys that look like integers back to int.
+        This reverses json.dump's coercion of dict keys to strings.
+        """
+        if isinstance(obj, dict):
+            new_obj = {}
+            for key, value in obj.items():
+                if isinstance(key, str) and key.lstrip('-').isdigit():
+                    new_key = int(key)
+                else:
+                    new_key = key
+                new_obj[new_key] = RosettaModel._convert_dict_keys_to_ints(value)
+            return new_obj
+        if isinstance(obj, list):
+            return [RosettaModel._convert_dict_keys_to_ints(v) for v in obj]
+        return obj
+    
+    
+    def save_projector_config(self, file_name):
+        with open(file_name, "w") as f:
+            json.dump(self.projector_dict, f)
+
+    
+    def load_projector_config(self, config_path):
+        if config_path.endswith(".json"):
+            loaded = RosettaModel.load_json(config_path)
+            self.projector_dict = RosettaModel._convert_dict_keys_to_ints(loaded)
+
+    def save_aggregator_config(self, file_name):
+        with open(file_name, "w") as f:
+            json.dump(self.aggregator_dict, f)
+
+    def load_aggregator_config(self, config_path):
+        if config_path.endswith(".json"):
+            loaded = RosettaModel.load_json(config_path)
+            self.aggregator_dict = RosettaModel._convert_dict_keys_to_ints(loaded)
+
+
+    def set_kv_cache_dict(self, source_model_idx, target_model_idx, cache):
+        if target_model_idx not in self.kv_cache_dict.keys():
+            self.kv_cache_dict[target_model_idx] = {}
+        if cache is None:
+            # Initialize with a DynamicCache instead of RosettaCache for now
+            self.kv_cache_dict[target_model_idx][source_model_idx] = DynamicCache() # noqa, maybe we should use RosettaCache here
+        else:
+            self.kv_cache_dict[target_model_idx][source_model_idx] = cache
+
+    def forward(
+        self,
+        kv_cache_index: Optional[List] = None,
+        input_ids: Optional[Union[torch.LongTensor, List[torch.LongTensor]]] = None,
+        attention_mask: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        # **kwargs: Unpack[KwargsForCausalLM],
+        identifier = -1,
+        subject = None,
+        *args,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """
+        Forward pass
+        KVCache index is a list of tensors with shape (B, sec_seq_len, 2), indicating the source and target kv cache model index
+
+        If input_ids is LongTensor, default to same input ids for different models
+        If input_ids is Tuple, default to different input ids for different models.
+
+        No Rosetta: (-1, 0)
+        """
+        
+        # noqa
+        self.kv_cache_dict = dict()
+
+        # Handle different input formats: if input_ids is a list, use per-model inputs
+        if isinstance(input_ids, list):
+            # Use list format: different input_ids and attention_mask for each model
+            base_input_ids = input_ids[self.base_model_idx] if input_ids is not None else None
+            base_attention_mask = attention_mask[self.base_model_idx] if attention_mask is not None else None
+            _, seqlen = base_input_ids.size() if base_input_ids is not None else (0, 0)
+        else:
+            # Use tensor format: same input_ids and attention_mask for all models (backward compatibility)
+            base_input_ids = input_ids
+            base_attention_mask = attention_mask
+            _, seqlen = input_ids.size() if input_ids is not None else (0, 0)
+
+        num_sections = len(kv_cache_index) if kv_cache_index is not None else 1
+
+        section_lengths = [kv_cache_index[i].shape[1] for i in range(num_sections)] if kv_cache_index is not None else [seqlen]
+        section_starts = [0]
+        for l in section_lengths:
+            section_starts.append(section_starts[-1] + l)
+        
+        curr_base_kv_cache = past_key_values
+
+        if seqlen > 1:
+            for i in range(num_sections):
+                start = section_starts[i]
+                end = section_starts[i + 1]
+                prefill_input_ids = base_input_ids[:, start:end] if base_input_ids is not None else None
+                prefill_attention_mask = base_attention_mask[:, :end] if base_attention_mask is not None else None
+                prefill_position_ids = position_ids[:, start:end] if position_ids is not None else None
+                prefill_labels = labels[:, start:end] if labels is not None else None
+
+                # calculate target model kvcache
+                output = self.model_list[self.base_model_idx].forward(
+                    input_ids=prefill_input_ids,
+                    attention_mask=prefill_attention_mask, 
+                    position_ids=prefill_position_ids,
+                    past_key_values=curr_base_kv_cache,
+                    labels=prefill_labels,
+                    use_cache=use_cache, 
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    *args,
+                    **kwargs
+                )
+
+                if self.base_model_idx not in self.kv_cache_dict:
+                    self.kv_cache_dict[self.base_model_idx] = {}
+                if self.base_model_idx not in self.kv_cache_dict[self.base_model_idx]:
+                    self.kv_cache_dict[self.base_model_idx][self.base_model_idx] = None
+                self.kv_cache_dict[self.base_model_idx][self.base_model_idx] = output.past_key_values
+
+                curr_base_kv_cache: DynamicCache = output.past_key_values
+                
+                # if i != num_sections - 1:
+                for source_model_idx in range(1, len(self.model_list)):
+                    if self.base_model_idx not in self.kv_cache_dict:
+                        self.kv_cache_dict[self.base_model_idx] = {}
+                    if source_model_idx not in self.kv_cache_dict[self.base_model_idx]:
+                        self.kv_cache_dict[self.base_model_idx][source_model_idx] = None
+
+                    # Get model-specific input_ids and attention_mask
+                    if isinstance(input_ids, list):
+                        source_input_ids = input_ids[source_model_idx]
+                        source_attention_mask = attention_mask[source_model_idx] if attention_mask is not None else None
+                        source_prefill_input_ids = source_input_ids[:, start:end] if source_input_ids is not None else None
+                        source_prefill_attention_mask = source_attention_mask[:, :end] if source_attention_mask is not None else None
+                    else:
+                        # Backward compatibility: use same input for all models
+                        source_prefill_input_ids = prefill_input_ids
+                        source_prefill_attention_mask = prefill_attention_mask
+
+                    curr_source_kv_cache = self.model_list[source_model_idx].forward(
+                        input_ids=source_prefill_input_ids,
+                        attention_mask=source_prefill_attention_mask,
+                        position_ids=prefill_position_ids,
+                        past_key_values=self.kv_cache_dict[self.base_model_idx][source_model_idx],
+                        use_cache=use_cache, 
+                        output_attentions=output_attentions,
+                        output_hidden_states=output_hidden_states,
+                        *args,
+                        **kwargs
+                    ).past_key_values
+                    self.kv_cache_dict[self.base_model_idx][source_model_idx] = curr_source_kv_cache
+
+                # calculate source model kvcache and apply projections
+                if self.base_model_idx in self.projector_dict:
+                    source_model_idx = kv_cache_index[i][0][0][0].item()  # Get the source model index from the kv_cache_index
+                    if source_model_idx != -1:
+                        for target_layer_idx, entry in self.projector_dict[self.base_model_idx][source_model_idx].items():
+                            base_key_cache, base_value_cache = curr_base_kv_cache[target_layer_idx]
+                            new_base_key_cache = base_key_cache[:, :, start:end, :]
+                            new_base_value_cache = base_value_cache[:, :, start:end, :]
+                            new_base_kv_cache = (new_base_key_cache, new_base_value_cache)
+
+                            pair_list = entry
+
+                            projected_kv_list = []
+                            source_kv_list = []
+                            for source_model_layer_idx, projector_idx in pair_list:
+                                source_key_cache, source_value_cache = self.kv_cache_dict[self.base_model_idx][source_model_idx][source_model_layer_idx]
+                                new_source_key_cache = source_key_cache[:, :, start:end, :]
+                                new_source_value_cache = source_value_cache[:, :, start:end, :]
+                                new_source_kv_cache = (new_source_key_cache, new_source_value_cache)
+                                projected_key, projected_value = self.projector_list[projector_idx].forward(
+                                    new_source_kv_cache, # tuple of (key, value), each of shape (B, N, H, D)
+                                    new_base_kv_cache
+                                )
+                                projected_kv_list.append((projected_key, projected_value))
+
+                                # --------------
+                                # save base and projected kv cache
+                                torch.save((projected_key, projected_value), f"oracle/projected_kv/{subject}_{identifier}_{i}.pt")
+                                torch.save(new_base_kv_cache, f"oracle/target_kv/{subject}_{identifier}_{i}.pt")
+                                # --------------
+                                source_kv_list.append(new_source_kv_cache)
+
+                            # Aggregate (fallback to first projector if no aggregator is available)
+                            use_aggregator = (
+                                len(projected_kv_list) > 1 and
+                                len(self.aggregator_list) > 0 and
+                                self.base_model_idx in self.aggregator_dict and
+                                source_model_idx in self.aggregator_dict[self.base_model_idx] and
+                                target_layer_idx in self.aggregator_dict[self.base_model_idx][source_model_idx]
+                            )
+
+                            if use_aggregator:
+                                aggregator_idx = self.aggregator_dict[self.base_model_idx][source_model_idx][target_layer_idx]
+                                agg_key, agg_value = self.aggregator_list[aggregator_idx].forward(
+                                    source_kv_list,
+                                    new_base_kv_cache,
+                                    projected_kv_list
+                                )
+                            else:
+                                # Fallback to first projector result when no aggregator is available
+                                agg_key, agg_value = projected_kv_list[0]
+
+                            # Update cache with aggregated result
+                            curr_base_kv_cache.key_cache[target_layer_idx][:, :, start:end, :] = agg_key
+                            curr_base_kv_cache.value_cache[target_layer_idx][:, :, start:end, :] = agg_value
+                        
+                        output.past_key_values = curr_base_kv_cache
+                                                                             
+        # use base model for decode phase
+        else:
+            # Handle list input format for decode phase as well
+            decode_input_ids = input_ids[self.base_model_idx] if isinstance(input_ids, list) else input_ids
+            decode_attention_mask = attention_mask[self.base_model_idx] if isinstance(attention_mask, list) else attention_mask
+            
+            output = self.model_list[self.base_model_idx].forward(
+                input_ids=decode_input_ids,
+                attention_mask=decode_attention_mask,
+                position_ids=position_ids,
+                past_key_values=curr_base_kv_cache,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                cache_position=cache_position,
+                *args,
+                **kwargs
+            )
+
+        return output
+
+    @torch.no_grad()
+    def generate(
+        self,
+        kv_cache_index,
+        input_ids,
+        max_new_tokens: Optional[int] = None,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        pad_token_id: Optional[int] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        do_sample: Optional[bool] = None,
+        return_dict_in_generate: Optional[bool] = None,
+        output_scores: Optional[bool] = None,
+        max_length: Optional[int] = None,
+        use_cache: bool = True,
+        *args,
+        **kwargs,
+    ):
+        """
+        New generation loop without using the base model's generate.
+        - Uses this module's forward for prefill and per-token decode.
+        - Samples tokens via rosetta.model.sampling.sample_token.
+        Returns a tensor of shape [batch, prompt_len + generated_len] for the base model stream.
+        """
+        # Derive number of tokens to generate
+        # If max_new_tokens not provided, infer from max_length
+        if isinstance(input_ids, list):
+            base_input_ids_for_len = input_ids[self.base_model_idx]
+        else:
+            base_input_ids_for_len = input_ids
+        prompt_len = base_input_ids_for_len.size(1)
+
+        # Default eos/pad from base model tokenizer/config if not provided
+        base_model = self.model_list[self.base_model_idx]
+        gen_cfg = getattr(base_model, "generation_config", None)
+        cfg_obj = gen_cfg if gen_cfg is not None else getattr(base_model, "config", None)
+        if eos_token_id is None and cfg_obj is not None:
+            eos_token_id = getattr(cfg_obj, "eos_token_id", None)
+        if pad_token_id is None and cfg_obj is not None:
+            pad_token_id = getattr(cfg_obj, "pad_token_id", None)
+        if pad_token_id is None and eos_token_id is not None:
+            pad_token_id = eos_token_id if isinstance(eos_token_id, int) else eos_token_id[0]
+
+        if max_new_tokens is None:
+            if max_length is not None:
+                if max_length <= prompt_len:
+                    max_new_tokens = 0
+                else:
+                    max_new_tokens = max_length - prompt_len
+            else:
+                raise ValueError("Provide max_new_tokens or max_length")
+        if max_new_tokens < 0:
+            raise ValueError("max_new_tokens must be non-negative")
+
+        # Resolve base inputs
+        if isinstance(input_ids, list):
+            base_input_ids = input_ids[self.base_model_idx]
+            base_attention_mask = attention_mask[self.base_model_idx] if attention_mask is not None else None
+        else:
+            base_input_ids = input_ids
+            base_attention_mask = attention_mask
+
+        if base_attention_mask is None:
+            base_attention_mask = torch.ones_like(base_input_ids, dtype=torch.long, device=base_input_ids.device)
+
+        batch_size = base_input_ids.size(0)
+
+        # Prefill to build caches and obtain initial logits
+        prefill_output = self.forward(
+            kv_cache_index=kv_cache_index,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            *args,
+            **kwargs,
+        )
+
+        current_past = prefill_output.past_key_values
+        all_input_ids = base_input_ids
+        current_attention_mask = base_attention_mask
+
+        # EOS handling setup
+        eos_set = None
+        if eos_token_id is not None:
+            eos_set = set(eos_token_id if isinstance(eos_token_id, list) else [eos_token_id])
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=all_input_ids.device)
+
+        # Start from last prefill logits
+        last_logits = prefill_output.logits[:, -1, :]
+
+        # Determine sampling mode
+        if do_sample is None:
+            do_sample = False
+        effective_temperature = temperature if do_sample else 0.0
+
+        # Optional scores collection
+        collect_scores = bool(return_dict_in_generate) and bool(output_scores)
+        scores = []
+
+        for _ in range(max_new_tokens):
+            if collect_scores:
+                scores.append(last_logits)
+            # Sample next token
+            next_token = sample_token(last_logits, temperature=effective_temperature, top_p=top_p, top_k=top_k)
+            if not isinstance(next_token, torch.Tensor):
+                next_token = torch.tensor([next_token], device=all_input_ids.device, dtype=torch.long).repeat(batch_size)
+
+            # Apply EOS logic
+            if eos_set is not None:
+                just_finished = torch.zeros_like(finished)
+                for eid in eos_set:
+                    just_finished |= (next_token == eid)
+                finished = finished | just_finished
+                if pad_token_id is not None:
+                    next_token = torch.where(
+                        finished,
+                        torch.tensor(pad_token_id, device=next_token.device, dtype=next_token.dtype),
+                        next_token,
+                    )
+
+            # Append sampled token
+            next_token_unsqueezed = next_token.unsqueeze(1)
+            all_input_ids = torch.cat([all_input_ids, next_token_unsqueezed], dim=1)
+            current_attention_mask = torch.cat(
+                [
+                    current_attention_mask,
+                    torch.ones((batch_size, 1), device=current_attention_mask.device, dtype=current_attention_mask.dtype),
+                ],
+                dim=1,
+            )
+
+            # Early stop if all sequences finished
+            if eos_set is not None and torch.all(finished):
+                break
+
+            # Decode one step using cached states; pass base-stream tensors
+            kv_cache_index = [torch.tensor([-1, 0], dtype=torch.long).repeat(1, 1).unsqueeze(0).to(all_input_ids.device)]
+
+            decode_output = self.forward(
+                kv_cache_index=kv_cache_index,
+                input_ids=next_token_unsqueezed,
+                attention_mask=current_attention_mask,
+                position_ids=None,
+                past_key_values=current_past,
+                use_cache=True,
+                *args,
+                **kwargs,
+            )
+            current_past = decode_output.past_key_values
+            last_logits = decode_output.logits[:, -1, :]
+
+        # Return style compatible with HF generate
+        if return_dict_in_generate:
+            if GreedySearchDecoderOnlyOutput is not None and SampleDecoderOnlyOutput is not None:
+                if do_sample:
+                    return SampleDecoderOnlyOutput(
+                        sequences=all_input_ids,
+                        scores=scores if collect_scores else None,
+                    )
+                else:
+                    return GreedySearchDecoderOnlyOutput(
+                        sequences=all_input_ids,
+                        scores=scores if collect_scores else None,
+                    )
+            # Fallback to generic ModelOutput
+            result = {"sequences": all_input_ids}
+            if collect_scores:
+                result["scores"] = scores
+            return ModelOutput(**result)
+        return all_input_ids

rosetta/model/projector.py

@@ -0,0 +1,1194 @@
+"""
+Projector nn module for the unified memory
+"""
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from transformers import Cache, DynamicCache
+from typing import Optional, Tuple, Literal, Union
+import copy
+import math
+
+from rosetta.utils.registry import register_model, get_projector_class, PROJECTOR_REGISTRY, capture_init_args, save_object, load_object
+
+class Projector(nn.Module):
+    """Base projector class for unified memory"""
+    
+    def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
+        """
+        Project and combine the source key-value tensors to the target key-value tensors
+        Args:
+            source_kv: Tuple of (key, value) tensors, each (..., D_s) where ... are arbitrary leading dimensions
+            target_kv: Tuple of (key, value) tensors, each (..., D_t) where ... are arbitrary leading dimensions
+        Returns:
+            Tuple of (key, value) tensors, each (..., D_t) with same leading dimensions as input
+        """
+        raise NotImplementedError("Subclasses must implement forward method")
+
+    def cache_project(self, source_kv_cache: Cache, target_kv_cache: Cache) -> Cache:
+        """
+        Project the source kv cache to the target kv cache
+        """
+        if not isinstance(source_kv_cache, DynamicCache) or not isinstance(target_kv_cache, DynamicCache):
+            raise ValueError("Only DynamicCache is supported")
+        
+        projected_cache = DynamicCache()
+        
+        # Process each layer
+        for layer_idx in range(len(source_kv_cache.key_cache)):
+            source_key = source_kv_cache.key_cache[layer_idx]  # (B, H, N, D_s)
+            source_value = source_kv_cache.value_cache[layer_idx]  # (B, H, N, D_s)
+            
+            # Get corresponding target tensors (for reference/combination)
+            if layer_idx < len(target_kv_cache.key_cache):
+                target_key = target_kv_cache.key_cache[layer_idx]  # (B, H, N, D_t)
+                target_value = target_kv_cache.value_cache[layer_idx]  # (B, H, N, D_t)
+            else:
+                # If target cache doesn't have this layer, create dummy tensors
+                B, H, N, D_s = source_key.shape
+                D_t = source_key.shape[-1]  # Assume same dimension for simplicity
+                target_key = torch.zeros(B, H, N, D_t, device=source_key.device, dtype=source_key.dtype)
+                target_value = torch.zeros(B, H, N, D_t, device=source_value.device, dtype=source_value.dtype)
+            
+            # Reshape for forward pass: DynamicCache format (B, H, N, D) -> projector format (B, N, H, D)
+            source_key_reshaped = source_key.transpose(1, 2)
+            source_value_reshaped = source_value.transpose(1, 2)
+            target_key_reshaped = target_key.transpose(1, 2)
+            target_value_reshaped = target_value.transpose(1, 2)
+            
+            # Project using forward method with tuple input/output
+            source_kv = (source_key_reshaped, source_value_reshaped)
+            target_kv = (target_key_reshaped, target_value_reshaped)
+            projected_key, projected_value = self.forward(source_kv, target_kv)
+            
+            # Reshape back: projector format (B, N, H, D) -> DynamicCache format (B, H, N, D)
+            projected_key = projected_key.transpose(1, 2)
+            projected_value = projected_value.transpose(1, 2)
+            
+            # Update cache
+            projected_cache.update(projected_key, projected_value, layer_idx)
+        
+        return projected_cache
+
+@register_model
+@capture_init_args
+class TrivialProjector(Projector):
+    """
+    Trivial projector that directly outputs the target key-value pairs without any modification.
+    This is useful as a baseline or when you want to effectively disable projection.
+    """
+    
+    def __init__(self, **kwargs):
+        """
+        Initialize the trivial projector.
+        
+        Args:
+            source_dim: Source dimension (ignored, kept for compatibility)
+            target_dim: Target dimension (ignored, kept for compatibility)
+            **kwargs: Additional arguments (ignored, kept for compatibility)
+        """
+        super().__init__()
+    
+    def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
+        """
+        Return the target key-value pairs unchanged, ignoring the source.
+        
+        Args:
+            source_kv: Tuple of (key, value) tensors (ignored)
+            target_kv: Tuple of (key, value) tensors to return unchanged
+        Returns:
+            The target key-value pairs unchanged
+        """
+        return target_kv
+
+@register_model
+@capture_init_args
+class ReplaceProjector(Projector):
+    """
+    Replacement projector that projects source key-value tensors to target dimension using MLP,
+    then replace target tensors using learnable weights.
+    """
+    def __init__(
+        self,
+        source_dim: int,
+        target_dim: int,
+        hidden_dim: int = 512,
+        num_layers: int = 2,
+        dropout: float = 0.1,
+        activation: str = "gelu",
+        use_layer_norm: bool = True,
+        init_weight: float = 0.1,
+        anneal_steps: int = 1360,
+        initial_temperature: float = 1.0,
+        final_temperature: float = 0.01,
+        scalar_temperature: float = 0.005,
+        # shared_key_projection: nn.Module = None,
+        # shared_value_projection: nn.Module = None,
+        dtype: torch.dtype = torch.float32
+    ):
+        super().__init__()
+
+        self.source_dim = source_dim
+        self.target_dim = target_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        
+        # Activation function
+        if activation.lower() == "gelu":
+            self.activation = nn.GELU()
+        elif activation.lower() == "relu":
+            self.activation = nn.ReLU()
+        elif activation.lower() == "silu":
+            self.activation = nn.SiLU()
+        else:
+            raise ValueError(f"Unsupported activation: {activation}")
+        
+        # Build separate MLP layers for key and value projection
+        self.key_projection = self._build_mlp(source_dim, hidden_dim, target_dim, num_layers, use_layer_norm, dropout, dtype)
+        self.value_projection = self._build_mlp(source_dim, hidden_dim, target_dim, num_layers, use_layer_norm, dropout, dtype)
+
+    
+    def _build_mlp(self, source_dim: int, hidden_dim: int, target_dim: int, num_layers: int, 
+                   use_layer_norm: bool, dropout: float, dtype: torch.dtype) -> nn.Sequential:
+        """Build a single MLP projection module"""
+        layers = []
+        
+        # Input projection
+        layers.append(nn.Linear(source_dim, hidden_dim, dtype=dtype))
+        if use_layer_norm:
+            layers.append(nn.LayerNorm(hidden_dim, dtype=dtype))
+        layers.append(copy.deepcopy(self.activation))
+        layers.append(nn.Dropout(dropout))
+        
+        # Hidden layers
+        for _ in range(num_layers - 2):
+            layers.append(nn.Linear(hidden_dim, hidden_dim, dtype=dtype))
+            if use_layer_norm:
+                layers.append(nn.LayerNorm(hidden_dim, dtype=dtype))
+            layers.append(copy.deepcopy(self.activation))
+            layers.append(nn.Dropout(dropout))
+        
+        # Output projection
+        if num_layers > 1:
+            layers.append(nn.Linear(hidden_dim, target_dim, dtype=dtype))
+        else:
+            # Single layer case
+            layers = [nn.Linear(source_dim, target_dim, dtype=dtype)]
+        
+        return nn.Sequential(*layers)
+
+
+    def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tensor]:
+        """
+        Project source key-value tensors to target dimension and add to target tensors with learnable weights
+        Args:
+            source_kv: Tuple of (key, value) tensors, each (..., D_s) where ... are arbitrary leading dimensions
+            target_kv: Tuple of (key, value) tensors, each (..., D_t) where ... are arbitrary leading dimensions
+        Returns:
+            Tuple of (key, value) tensors, each (..., D_t) with same leading dimensions as input
+        """
+
+        source_key, source_value = source_kv
+        target_key, target_value = target_kv
+        
+        # Get shapes - assuming format is (B, H, N, D) where H is num_heads, N is seq_len, D is head_dim
+        source_shape = source_key.shape  # (B, H_s, N, D_s)
+        target_shape = target_key.shape  # (B, H_t, N, D_t)
+        
+        # Extract dimensions
+        batch_size, source_num_heads, seq_len, source_head_dim = source_shape
+        _, target_num_heads, _, target_head_dim = target_shape
+        
+        # Reshape source: merge num_heads and head_dim for projection
+        # (B, H_s, N, D_s) -> (B, N, H_s * D_s)
+        source_key_reshaped = source_key.transpose(1, 2)  # (B, N, H_s, D_s)
+        source_value_reshaped = source_value.transpose(1, 2)  # (B, N, H_s, D_s)
+        
+        source_key_flat = source_key_reshaped.contiguous().view(batch_size, seq_len, source_num_heads * source_head_dim)
+        source_value_flat = source_value_reshaped.contiguous().view(batch_size, seq_len, source_num_heads * source_head_dim)
+        
+        # Project source tensors from (H_s * D_s) to (H_t * D_t)
+        projected_key_flat = self.key_projection(source_key_flat)  # (B, N, H_t * D_t)
+        projected_value_flat = self.value_projection(source_value_flat)  # (B, N, H_t * D_t)
+        
+        # Reshape projected tensors back to target format
+        # (B, N, H_t * D_t) -> (B, N, H_t, D_t) -> (B, H_t, N, D_t)
+        projected_key_reshaped = projected_key_flat.view(batch_size, seq_len, target_num_heads, target_head_dim)
+        projected_value_reshaped = projected_value_flat.view(batch_size, seq_len, target_num_heads, target_head_dim)
+        
+        projected_key = projected_key_reshaped.transpose(1, 2)  # (B, H_t, N, D_t)
+        projected_value = projected_value_reshaped.transpose(1, 2)  # (B, H_t, N, D_t)
+        
+        return (projected_key, projected_value)
+
+class ModernMLP(nn.Module):
+    """
+    Modern MLP with residual connections, layer normalization, and configurable architecture.
+    """
+    
+    def __init__(
+        self,
+        input_dim: int,
+        output_dim: int,
+        hidden_dim: int = 512,
+        num_layers: int = 2,
+        activation: str = "gelu",
+        use_layer_norm: bool = True,
+        use_residual: bool = True,
+        dropout: float = 0.1,
+        use_swiglu: bool = False,
+        dtype: torch.dtype = torch.float32
+    ):
+        super().__init__()
+        
+        self.input_dim = input_dim
+        self.output_dim = output_dim
+        self.hidden_dim = hidden_dim
+        self.num_layers = num_layers
+        self.use_residual = use_residual and (input_dim == output_dim)
+        self.use_swiglu = use_swiglu
+        
+        # Activation function
+        if activation.lower() == "gelu":
+            self.activation = nn.GELU()
+        elif activation.lower() == "relu":
+            self.activation = nn.ReLU()
+        elif activation.lower() == "silu":
+            self.activation = nn.SiLU()
+        else:
+            raise ValueError(f"Unsupported activation: {activation}")
+        
+        # Build layers
+        self.layers = nn.ModuleList()
+        
+        for i in range(num_layers):
+            layer_input_dim = input_dim if i == 0 else hidden_dim
+            layer_output_dim = output_dim if i == num_layers - 1 else hidden_dim
+            
+            if self.use_swiglu and i < num_layers - 1:  # Don't use SwiGLU on output layer
+                layer = SwiGLUBlock(layer_input_dim, layer_output_dim, dtype=dtype)
+            else:
+                layer = nn.Linear(layer_input_dim, layer_output_dim, dtype=dtype)
+            
+            self.layers.append(layer)
+            
+            # Add layer norm after each layer except the last one
+            if use_layer_norm and i < num_layers - 1:
+                self.layers.append(nn.LayerNorm(layer_output_dim, dtype=dtype))
+            
+            # Add activation after each layer except the last one
+            if i < num_layers - 1 and not self.use_swiglu:
+                self.layers.append(copy.deepcopy(self.activation))
+            
+            # Add dropout after activation
+            if dropout > 0 and i < num_layers - 1:
+                self.layers.append(nn.Dropout(dropout))
+        
+        # Residual projection if dimensions don't match
+        if self.use_residual and input_dim != output_dim:
+            self.residual_proj = nn.Linear(input_dim, output_dim, dtype=dtype)
+        else:
+            self.residual_proj = None
+    
+    def forward(self, x: Tensor) -> Tensor:
+        """Forward pass with optional residual connection."""
+        residual = x
+        
+        for layer in self.layers:
+            x = layer(x)
+        
+        # Add residual connection
+        if self.use_residual:
+            if self.residual_proj is not None:
+                residual = self.residual_proj(residual)
+            x = x + residual
+        
+        return x
+
+
+class SwiGLUBlock(nn.Module):
+    """SwiGLU activation block for modern transformer architectures."""
+    
+    def __init__(self, input_dim: int, output_dim: int, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.gate_proj = nn.Linear(input_dim, output_dim, dtype=dtype)
+        self.up_proj = nn.Linear(input_dim, output_dim, dtype=dtype)
+        self.activation = nn.SiLU()
+    
+    def forward(self, x: Tensor) -> Tensor:
+        gate = self.activation(self.gate_proj(x))
+        up = self.up_proj(x)
+        return gate * up
+
+
+@register_model
+@capture_init_args
+class AllInOneProjector(Projector):
+    """
+    Unified projector that consolidates all projection functionalities with modern patterns.
+    
+    Features:
+    1. Gate logit granularity: scalar, token-wise, head-wise, or value-wise
+    2. (DISABLED) Selector logit granularity: scalar, token-wise, head-wise, or value-wise
+    3. Key/Value weight granularity: scalar, token-wise, head-wise, or value-wise
+    4. Input-dependent gates and weights via MLP or parameters
+    5. Optional concatenation with combiner networks
+    6. Modern MLP architecture with residual connections and SwiGLU
+    7. Configurable target preservation: choose between traditional blending or simplified projection
+    8. Optional adding of target (self) signal to outputs via add_self
+    
+    Target Preservation Modes:
+    - preserve_target_weight=True (default): output = (1-weight)*target + gate*selector*weight*projected
+    - preserve_target_weight=False: output = target + gate*selector*weight*projected (no weight coefficient on target)
+    
+    Note: Selector functionality has been disabled/commented out.
+    """
+    
+    def __init__(
+        self,
+        source_dim: int,
+        target_dim: int,
+        source_num_heads: int = 1,
+        target_num_heads: int = 1,
+        hidden_dim: int = 512,
+        num_layers: int = 2,
+        dropout: float = 0.1,
+        activation: str = "gelu",
+        use_layer_norm: bool = True,
+        use_residual: bool = True,
+        use_swiglu: bool = False,
+        
+        # Gate configuration
+        gate_granularity: Literal["scalar", "token", "head", "head_merged", "value"] = "scalar",
+        gate_depends_on_input: bool = False,
+        gate_input_features: Optional[str] = "target_key",  # "target_key", "target_value", "both", "target_projected_key", "target_projected_value", "target_projected_both"
+        gate_init_value: float = 0.0,
+        
+        # Weight configuration
+        weight_granularity: Literal["scalar", "token", "head", "head_merged", "value"] = "scalar",
+        weight_depends_on_input: bool = False,
+        weight_input_features: Optional[str] = "target_key",  # "target_key", "target_value", "both", "target_projected_key", "target_projected_value", "target_projected_both"
+        weight_init_value: float = 0.0,
+        
+        # Target preservation configuration
+        preserve_target_weight: bool = True,  # If False, target won't be multiplied by (1 - normalized_weight)
+        add_self: bool = True,  # If False, target (self) won't be added to outputs
+        
+        # Concat configuration
+        use_concat: bool = False,
+        # combiner_hidden_dim: int = 128,
+        weight_hidden_dim: int = 1024,
+        
+        # Temperature and gumbel
+        use_gumbel: bool = True,
+        initial_temperature: float = 1.0,
+        final_temperature: float = 0.01,
+        anneal_steps: int = 1360,
+        scalar_temperature: float = 0.005,
+        
+        # Sequence length configuration
+        max_sequence_length: int = 8192,  # Maximum sequence length for token-level parameters
+
+        pos_emb: bool = False,
+
+        dtype: torch.dtype = torch.float32
+    ):
+        super().__init__()
+        
+        self.source_dim = source_dim
+        self.target_dim = target_dim
+        self.source_num_heads = source_num_heads
+        self.target_num_heads = target_num_heads
+        self.hidden_dim = hidden_dim
+        self.weight_hidden_dim = weight_hidden_dim
+        self.max_sequence_length = max_sequence_length
+        
+        # Configuration
+        self.gate_granularity = gate_granularity
+        self.gate_depends_on_input = gate_depends_on_input
+        self.gate_input_features = gate_input_features
+        self.weight_granularity = weight_granularity
+        self.weight_depends_on_input = weight_depends_on_input
+        self.weight_input_features = weight_input_features
+        self.preserve_target_weight = preserve_target_weight
+        self.add_self = add_self
+        self.use_concat = use_concat
+        self.use_gumbel = use_gumbel
+        self.scalar_temperature = scalar_temperature
+
+        # Temperature annealing for gate only (removed selector temperature)
+        self.register_buffer("gate_temperature", torch.tensor(initial_temperature, dtype=dtype))
+        self.initial_temperature = initial_temperature
+        self.final_temperature = final_temperature
+        self.anneal_steps = anneal_steps
+        
+        # Build projection networks
+        self.key_projection = self._build_projection_mlp(
+            source_dim * source_num_heads, 
+            target_dim * target_num_heads,
+            hidden_dim, num_layers, activation, use_layer_norm, 
+            use_residual, dropout, use_swiglu, dtype
+        )
+        self.value_projection = self._build_projection_mlp(
+            source_dim * source_num_heads,
+            target_dim * target_num_heads,
+            hidden_dim, num_layers, activation, use_layer_norm,
+            use_residual, dropout, use_swiglu, dtype
+        )
+        
+        # Build gate components
+        self._build_gate_components(dtype)
+        
+        # Build weight components  
+        self._build_weight_components(weight_init_value, dtype)
+
+        # Build concat components if needed
+        if self.use_concat:
+            in_dim = target_dim * target_num_heads * 2
+            out_dim = target_dim * target_num_heads
+            self.key_combiner = nn.Linear(in_dim, out_dim, dtype=dtype)
+            self.value_combiner = nn.Linear(in_dim, out_dim, dtype=dtype)
+        
+    def _build_projection_mlp(
+        self, input_dim: int, output_dim: int, hidden_dim: int, 
+        num_layers: int, activation: str, use_layer_norm: bool,
+        use_residual: bool, dropout: float, use_swiglu: bool, dtype: torch.dtype
+    ) -> ModernMLP:
+        """Build modern MLP for projection."""
+        return ModernMLP(
+            input_dim=input_dim,
+            output_dim=output_dim,
+            hidden_dim=hidden_dim,
+            num_layers=num_layers,
+            activation=activation,
+            use_layer_norm=use_layer_norm,
+            use_residual=use_residual,
+            dropout=dropout,
+            use_swiglu=use_swiglu,
+            dtype=dtype
+        )
+    
+    def _build_gate_components(self, dtype: torch.dtype):
+        """Build gate logit components based on configuration."""
+        if not self.gate_depends_on_input:
+            # Parameter-based gate
+            gate_shape = self._get_parameter_shape(self.gate_granularity)
+            self.gate_logit = nn.Parameter(torch.zeros(gate_shape, dtype=dtype))
+        else:
+            # Input-dependent gate via MLP
+            input_dim = self._get_gate_input_dim()
+            output_dim = self._get_gate_output_dim()
+            
+            self.gate_generator = ModernMLP(
+                input_dim=input_dim,
+                output_dim=output_dim,
+                hidden_dim=self.hidden_dim,
+                num_layers=2,
+                activation="gelu",
+                use_layer_norm=True,
+                use_residual=False,
+                dropout=0.1,
+                dtype=dtype
+            )
+    
+    def _build_weight_components(self, weight_init_value: float, dtype: torch.dtype):
+        """Build weight components based on configuration."""
+        if not self.weight_depends_on_input:
+            # Parameter-based weights
+            weight_shape = self._get_parameter_shape(self.weight_granularity)
+            self.key_weight = nn.Parameter(torch.full(weight_shape, weight_init_value, dtype=dtype))
+            self.value_weight = nn.Parameter(torch.full(weight_shape, weight_init_value, dtype=dtype))
+        else:
+            # Input-dependent weights via MLP
+            input_dim = self._get_weight_input_dim()
+            output_dim = self._get_weight_output_dim()
+            
+            # Shared hidden layer for efficiency
+            self.weight_hidden = ModernMLP(
+                input_dim=input_dim,
+                output_dim=self.weight_hidden_dim,
+                hidden_dim=self.weight_hidden_dim,
+                num_layers=2,
+                activation="gelu",
+                use_layer_norm=True,
+                use_residual=False,
+                dropout=0.1,
+                dtype=dtype
+            )
+            
+            # Separate heads for key and value weights
+            self.key_weight_head = nn.Linear(self.weight_hidden_dim, output_dim, dtype=dtype)
+            self.value_weight_head = nn.Linear(self.weight_hidden_dim, output_dim, dtype=dtype)
+    
+    def _get_parameter_shape(self, granularity: str) -> tuple:
+        """Get parameter shape based on granularity."""
+        if granularity == "scalar":
+            return ()  # Scalar
+        elif granularity == "token":
+            return (self.max_sequence_length,)  # Token-level parameters with max sequence length
+        elif granularity == "head":
+            return (self.max_sequence_length, self.target_num_heads)  # Token and head level parameters
+        elif granularity == "head_merged":
+            return (self.max_sequence_length, self.target_num_heads)  # Token and head level parameters
+        elif granularity == "value":
+            return (self.max_sequence_length, self.target_num_heads, self.target_dim)  # Token, head and value level parameters
+        else:
+            raise ValueError(f"Invalid granularity: {granularity}")
+    
+    def _get_gate_input_dim(self) -> int:
+        """Get input dimension for gate generator."""
+        base_dim = 0
+        if self.gate_input_features == "target_key":
+            base_dim = self.target_dim
+        elif self.gate_input_features == "target_value":
+            base_dim = self.target_dim
+        elif self.gate_input_features == "both":
+            base_dim = self.target_dim * 2
+        elif self.gate_input_features == "target_projected_key":
+            base_dim = self.target_dim * 2  # target_key + projected_key
+        elif self.gate_input_features == "target_projected_value":
+            base_dim = self.target_dim * 2  # target_value + projected_value
+        elif self.gate_input_features == "target_projected_both":
+            base_dim = self.target_dim * 4  # target_key + target_value + projected_key + projected_value
+        else:
+            raise ValueError(f"Invalid gate input features: {self.gate_input_features}")
+        
+        # Adjust for granularity processing strategy
+        if self.gate_granularity == "scalar":
+            # Scalar: process aggregated features across all heads
+            return base_dim  # Use pooled features
+        elif self.gate_granularity == "token":
+            # Token: process merged head dimensions 
+            return base_dim * self.target_num_heads  # Flatten (H, D) to (H*D)
+        elif self.gate_granularity == "head_merged":
+            # Head-merged: similar to token granularity, merge H and D
+            return base_dim * self.target_num_heads  # (B, N, H*D)
+        elif self.gate_granularity == "head":
+            # Head-local: per head processing, do not merge heads
+            return base_dim  # (B, H, N, D)
+        else:  # value
+            # Value: process per-head features
+            return base_dim  # Keep per-head processing (B, H, N, D)
+    
+    def _get_gate_output_dim(self) -> int:
+        """Get output dimension for gate generator."""
+        if self.gate_granularity == "scalar":
+            return 1
+        elif self.gate_granularity == "token":
+            return 1  # Per token
+        elif self.gate_granularity == "head_merged":
+            # Per token per head after merge: output one value per head
+            return self.target_num_heads
+        elif self.gate_granularity == "head":
+            # Per token per head: scalar per head
+            return 1
+        elif self.gate_granularity == "value":
+            return self.target_dim  # Per token per head per value (but processed per-head, so output D per head)
+        else:
+            raise ValueError(f"Invalid gate granularity: {self.gate_granularity}")
+    
+    def _get_weight_input_dim(self) -> int:
+        """Get input dimension for weight generator."""
+        base_dim = 0
+        if self.weight_input_features == "target_key":
+            base_dim = self.target_dim
+        elif self.weight_input_features == "target_value":
+            base_dim = self.target_dim
+        elif self.weight_input_features == "both":
+            base_dim = self.target_dim * 2
+        elif self.weight_input_features == "target_projected_key":
+            base_dim = self.target_dim * 2  # target_key + projected_key
+        elif self.weight_input_features == "target_projected_value":
+            base_dim = self.target_dim * 2  # target_value + projected_value
+        elif self.weight_input_features == "target_projected_both":
+            base_dim = self.target_dim * 4  # target_key + target_value + projected_key + projected_value
+        else:
+            raise ValueError(f"Invalid weight input features: {self.weight_input_features}")
+        
+        # Adjust for granularity processing strategy
+        if self.weight_granularity == "scalar":
+            # Scalar: process aggregated features across all heads
+            return base_dim  # Use pooled features
+        elif self.weight_granularity == "token":
+            # Token: process merged head dimensions 
+            return base_dim * self.target_num_heads  # Flatten (H, D) to (H*D)
+        elif self.weight_granularity == "head_merged":
+            # Head-merged: similar to token granularity, merge H and D
+            return base_dim * self.target_num_heads  # (B, N, H*D)
+        elif self.weight_granularity == "head":
+            # Head-local: per head processing, do not merge heads
+            return base_dim  # (B, H, N, D)
+        else:  # value
+            # Value: process per-head features
+            return base_dim  # Keep per-head processing (B, H, N, D)
+    
+    def _get_weight_output_dim(self) -> int:
+        """Get output dimension for weight generator."""
+        if self.weight_granularity == "scalar":
+            return 1
+        elif self.weight_granularity == "token":
+            return 1  # Per token
+        elif self.weight_granularity == "head_merged":
+            # Per token per head after merge: output one value per head
+            return self.target_num_heads
+        elif self.weight_granularity == "head":
+            # Per token per head: scalar per head
+            return 1
+        elif self.weight_granularity == "value":
+            return self.target_dim  # Per token per head per value (but processed per-head, so output D per head)
+        else:
+            raise ValueError(f"Invalid weight granularity: {self.weight_granularity}")
+    
+    def _generate_gates(self, target_key: Tensor, target_value: Tensor, projected_key: Tensor = None, projected_value: Tensor = None) -> Tensor:
+        """Generate gate logits based on configuration."""
+        if not self.gate_depends_on_input:
+            # Use parameter-based gate
+            return self.gate_logit
+        else:
+            # Generate input-dependent gate
+            # First, prepare the base input features
+            if self.gate_input_features == "target_key":
+                base_input = target_key
+            elif self.gate_input_features == "target_value":
+                base_input = target_value
+            elif self.gate_input_features == "both":
+                base_input = torch.cat([target_key, target_value], dim=-1)
+            elif self.gate_input_features == "target_projected_key":
+                if projected_key is None:
+                    raise ValueError("projected_key is required for target_projected_key input features")
+                base_input = torch.cat([target_key, projected_key], dim=-1)
+            elif self.gate_input_features == "target_projected_value":
+                if projected_value is None:
+                    raise ValueError("projected_value is required for target_projected_value input features")
+                base_input = torch.cat([target_value, projected_value], dim=-1)
+            elif self.gate_input_features == "target_projected_both":
+                if projected_key is None or projected_value is None:
+                    raise ValueError("Both projected_key and projected_value are required for target_projected_both input features")
+                base_input = torch.cat([target_key, target_value, projected_key, projected_value], dim=-1)
+            
+            # Now process based on granularity
+            # base_input shape: (B, H, N, D_input)
+            B, H, N, D_input = base_input.shape
+            
+            if self.gate_granularity == "scalar":
+                # For scalar granularity, aggregate all dimensions: (B, H, N, D_input) -> (B, D_input)
+                gate_input = base_input.mean(dim=(1, 2))  # Average over heads and tokens
+            elif self.gate_granularity == "token":
+                # For token granularity, merge H and D_input dimensions: (B, H, N, D_input) -> (B, N, H*D_input)
+                gate_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input)
+            elif self.gate_granularity == "head_merged":
+                # For head granularity, merge H and D like token: (B, H, N, D_in) -> (B, N, H*D_in)
+                gate_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input)
+            elif self.gate_granularity == "head":
+                # For head granularity, keep per-head processing: (B, H, N, D_input)
+                gate_input = base_input
+            elif self.gate_granularity == "value":
+                # For value granularity, keep per-head processing: (B, H, N, D_input)
+                gate_input = base_input
+            
+            return self.gate_generator(gate_input)
+    
+    def _generate_weights(self, target_key: Tensor, target_value: Tensor, projected_key: Tensor = None, projected_value: Tensor = None) -> Tuple[Tensor, Tensor]:
+        """Generate weights based on configuration."""
+        if not self.weight_depends_on_input:
+            # Use parameter-based weights
+            return self.key_weight, self.value_weight
+        else:
+            # Generate input-dependent weights
+            # First, prepare the base input features
+            if self.weight_input_features == "target_key":
+                base_input = target_key
+            elif self.weight_input_features == "target_value":
+                base_input = target_value
+            elif self.weight_input_features == "both":
+                base_input = torch.cat([target_key, target_value], dim=-1)
+            elif self.weight_input_features == "target_projected_key":
+                if projected_key is None:
+                    raise ValueError("projected_key is required for target_projected_key input features")
+                base_input = torch.cat([target_key, projected_key], dim=-1)
+            elif self.weight_input_features == "target_projected_value":
+                if projected_value is None:
+                    raise ValueError("projected_value is required for target_projected_value input features")
+                base_input = torch.cat([target_value, projected_value], dim=-1)
+            elif self.weight_input_features == "target_projected_both":
+                if projected_key is None or projected_value is None:
+                    raise ValueError("Both projected_key and projected_value are required for target_projected_both input features")
+                base_input = torch.cat([target_key, target_value, projected_key, projected_value], dim=-1)
+            
+            # Now process based on granularity
+            # base_input shape: (B, H, N, D_input)
+            B, H, N, D_input = base_input.shape
+            
+            if self.weight_granularity == "scalar":
+                # For scalar granularity, aggregate all dimensions: (B, H, N, D_input) -> (B, D_input)
+                weight_input = base_input.mean(dim=(1, 2))  # Average over heads and tokens
+            elif self.weight_granularity == "token":
+                # For token granularity, merge H and D_input dimensions: (B, H, N, D_input) -> (B, N, H*D_input)
+                weight_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input)
+            elif self.weight_granularity == "head_merged":
+                # For head granularity, merge H and D like token: (B, H, N, D_in) -> (B, N, H*D_in)
+                weight_input = base_input.transpose(1, 2).contiguous().view(B, N, H * D_input)
+            elif self.weight_granularity == "head":
+                # For head granularity, keep per-head processing: (B, H, N, D_input)
+                weight_input = base_input
+            elif self.weight_granularity == "value":
+                # For value granularity, keep per-head processing: (B, H, N, D_input)
+                weight_input = base_input
+    
+            weight_hidden = self.weight_hidden(weight_input)
+            key_weight = self.key_weight_head(weight_hidden)
+            value_weight = self.value_weight_head(weight_hidden)
+
+            return key_weight, value_weight
+    
+    def _apply_gumbel_sigmoid(self, gate_logit: Tensor) -> Tensor:
+        """Apply Gumbel sigmoid trick for training."""
+        if self.training and self.use_gumbel:
+            gumbel_noise = self._sample_gumbel(gate_logit.shape, gate_logit.device, gate_logit.dtype)
+            return torch.sigmoid((gate_logit + gumbel_noise) / self.gate_temperature)
+        else:
+            return (gate_logit > 0).float()
+    
+    @staticmethod
+    def _sample_gumbel(shape: tuple, device: torch.device, dtype: torch.dtype, eps: float = 1e-20) -> Tensor:
+        """Sample from Gumbel distribution."""
+        u = torch.rand(shape, device=device, dtype=dtype)
+        return -torch.log(-torch.log(u + eps) + eps)
+    
+    def _reshape_for_granularity(self, tensor: Tensor, granularity: str, target_shape: tuple) -> Tensor:
+        """Reshape tensor to match target shape based on granularity."""
+        B, H, N, D = target_shape
+        
+        if granularity == "scalar":
+            # Scalar -> (B, H, N, D)
+            return tensor.view(1, 1, 1, 1).expand(B, H, N, D)
+        elif granularity == "token":
+            # (max_seq_len,) -> (B, H, N, D) - slice to actual sequence length
+            token_params = tensor[:N]  # Take first N tokens
+            return token_params.view(1, 1, N, 1).expand(B, H, N, D)
+        elif granularity == "head":
+            # (max_seq_len, H) -> (B, H, N, D) - slice to actual sequence length, each token each head independent
+            head_params = tensor[:N, :]  # Take first N tokens, all heads: (N, H)
+            return head_params.view(1, N, H, 1).transpose(1, 2).expand(B, H, N, D)  # (1, N, H, 1) -> (1, H, N, 1) -> (B, H, N, D)
+        elif granularity == "head_merged":
+            raise NotImplementedError
+        elif granularity == "value":
+            # (max_seq_len, H, D) -> (B, H, N, D) - slice to actual sequence length, each token each head each value independent
+            value_params = tensor[:N, :, :]  # Take first N tokens: (N, H, D)
+            return value_params.view(1, N, H, D).transpose(1, 2).expand(B, H, N, D)  # (1, N, H, D) -> (1, H, N, D) -> (B, H, N, D)
+        else:
+            raise ValueError(f"Invalid granularity: {granularity}")
+    
+    def update_temperature(self, step: int):
+        """Update temperature using exponential annealing schedule for gate only."""
+        # Update gate temperature
+        gate_ratio = min(step / self.anneal_steps, 1.0)
+        gate_temp = self.initial_temperature * (self.final_temperature / self.initial_temperature) ** gate_ratio
+        self.gate_temperature.fill_(gate_temp)
+    
+    
+    def forward(self, source_kv: Tuple[Tensor, Tensor], target_kv: Tuple[Tensor, Tensor], position_ids: Optional[Tensor] = None, max_pos: Optional[Tensor] = None) -> Tuple[Tensor, Tensor]:
+        """
+        Forward pass with unified projection logic.
+        
+        Args:
+            source_kv: Tuple of (key, value) tensors, each (B, H_s, N, D_s)
+            target_kv: Tuple of (key, value) tensors, each (B, H_t, N, D_t)
+            position_ids: Position ids tensor (B, N), optional, required if pos_emb=True
+        Returns:
+            Tuple of (key, value) tensors, each (B, H_t, N, D_t)
+        """
+        source_key, source_value = source_kv
+        target_key, target_value = target_kv
+        
+        # Get shapes
+        B, H_s, N, D_s = source_key.shape
+        _, H_t, _, D_t = target_key.shape
+        
+        # Reshape for projection: (B, H, N, D) -> (B, N, H*D)
+        source_key_flat = source_key.transpose(1, 2).contiguous().view(B, N, H_s * D_s)
+        source_value_flat = source_value.transpose(1, 2).contiguous().view(B, N, H_s * D_s)
+        
+        # Project source to target dimension
+        projected_key_flat = self.key_projection(source_key_flat)  # (B, N, H_t * D_t)
+        projected_value_flat = self.value_projection(source_value_flat)  # (B, N, H_t * D_t)
+        
+        # Handle concatenation if enabled
+        if self.use_concat:
+            target_key_flat = target_key.transpose(1, 2).contiguous().view(B, N, H_t * D_t)
+            target_value_flat = target_value.transpose(1, 2).contiguous().view(B, N, H_t * D_t)
+            
+            # Concatenate and combine
+            combined_key = torch.cat([projected_key_flat, target_key_flat], dim=-1)
+            combined_value = torch.cat([projected_value_flat, target_value_flat], dim=-1)
+            
+            final_projected_key_flat = self.key_combiner(combined_key)
+            final_projected_value_flat = self.value_combiner(combined_value)
+        else:
+            final_projected_key_flat = projected_key_flat
+            final_projected_value_flat = projected_value_flat
+        
+        # Reshape back: (B, N, H_t * D_t) -> (B, H_t, N, D_t)
+        projected_key = final_projected_key_flat.view(B, N, H_t, D_t).transpose(1, 2)
+        projected_value = final_projected_value_flat.view(B, N, H_t, D_t).transpose(1, 2)
+        
+        # Generate gates, selectors and weights (may need projected tensors for input features)
+        needs_projected_for_gate = self.gate_depends_on_input and self.gate_input_features in [
+            "target_projected_key", "target_projected_value", "target_projected_both"
+        ]
+        needs_projected_for_weight = self.weight_depends_on_input and self.weight_input_features in [
+            "target_projected_key", "target_projected_value", "target_projected_both"
+        ]
+        
+        if needs_projected_for_gate or needs_projected_for_weight:
+            gate_logit = self._generate_gates(target_key, target_value, projected_key, projected_value)
+            key_weight, value_weight = self._generate_weights(target_key, target_value, projected_key, projected_value)
+        else:
+            gate_logit = self._generate_gates(target_key, target_value)
+            key_weight, value_weight = self._generate_weights(target_key, target_value)
+        
+        # Reshape gates and weights to match target shape
+        target_shape = (B, H_t, N, D_t)
+        if self.gate_depends_on_input:
+            # Reshape based on gate granularity - all preserve token dimension N
+            if self.gate_granularity == "scalar":
+                # For scalar, gate_logit is already (B, 1) from MLP, just expand
+                gate_logit = gate_logit.view(B, 1, 1, 1).expand(target_shape)
+            elif self.gate_granularity == "token":
+                gate_logit = gate_logit.unsqueeze(1).unsqueeze(-1).expand(target_shape)  # (B, N, 1) -> (B, H, N, D)
+            elif self.gate_granularity == "head_merged":
+                # (B, N, H) -> (B, H, N, D) - per token per head, broadcast over D
+                gate_logit = gate_logit.permute(0, 2, 1).unsqueeze(-1).expand(B, H_t, N, D_t)
+            elif self.gate_granularity == "head":
+                # (B, H, N, 1) -> (B, H, N, D) - per token per head scalar, broadcast over D
+                gate_logit = gate_logit.expand(B, H_t, N, D_t)
+            elif self.gate_granularity == "value":
+                # (B, H, N, D) -> (B, H, N, D) - each token each head each value has one value
+                pass  # Already in correct shape
+        else:
+            gate_logit = self._reshape_for_granularity(gate_logit, self.gate_granularity, target_shape)
+
+        if self.weight_depends_on_input:
+            # Reshape weights based on granularity - all preserve token dimension N
+            if self.weight_granularity == "scalar":
+                # For scalar, weights are already (B, 1) from MLP, just expand
+                key_weight = key_weight.view(B, 1, 1, 1).expand(target_shape)
+                value_weight = value_weight.view(B, 1, 1, 1).expand(target_shape)
+            elif self.weight_granularity == "token":
+                key_weight = key_weight.unsqueeze(1).expand(target_shape)  # (B, N, 1) -> (B, H, N, D)
+                value_weight = value_weight.unsqueeze(1).expand(target_shape)
+            elif self.weight_granularity == "head_merged":
+                # (B, N, H) -> (B, H, N, D) - per token per head, broadcast over D
+                key_weight = key_weight.permute(0, 2, 1).unsqueeze(-1).expand(B, H_t, N, D_t)
+                value_weight = value_weight.permute(0, 2, 1).unsqueeze(-1).expand(B, H_t, N, D_t)
+            elif self.weight_granularity == "head":
+                # (B, H, N, 1) -> (B, H, N, D) - per token per head scalar, broadcast over D
+                key_weight = key_weight.expand(B, H_t, N, D_t)
+                value_weight = value_weight.expand(B, H_t, N, D_t)
+            elif self.weight_granularity == "value":
+                # (B, H, N, D) -> (B, H, N, D) - each token each head each value has one value
+                pass  # Already in correct shape
+        else:
+            key_weight = self._reshape_for_granularity(key_weight, self.weight_granularity, target_shape)
+            value_weight = self._reshape_for_granularity(value_weight, self.weight_granularity, target_shape)
+        
+        # Apply gating and selection
+        gate = self._apply_gumbel_sigmoid(gate_logit)
+        
+        # Normalize weights using dynamic temperature
+        normalized_key_weight = torch.sigmoid(key_weight / self.scalar_temperature)
+        normalized_value_weight = torch.sigmoid(value_weight / self.scalar_temperature)
+        
+        # Final combination
+        # Compute projected contribution (always present)
+        projected_key_term = gate * normalized_key_weight * projected_key
+        projected_value_term = gate * normalized_value_weight * projected_value
+
+        # Compute target (self) contribution depending on flags
+        if self.add_self:
+            if self.preserve_target_weight:
+                target_key_term = (1 - normalized_key_weight) * target_key
+                target_value_term = (1 - normalized_value_weight) * target_value
+            else:
+                target_key_term = target_key
+                target_value_term = target_value
+        else:
+            target_key_term = torch.zeros_like(target_key)
+            target_value_term = torch.zeros_like(target_value)
+
+        # Final outputs
+        output_key = target_key_term + projected_key_term
+        output_value = target_value_term + projected_value_term
+        
+        return (output_key, output_value)
+
+class QwenStyleLayer(nn.Module):
+    """
+    One Qwen3-style MLP sublayer:
+      y = x + Dropout( down( SiLU(gate(LN(x))) * up(LN(x)) ) )
+    - Pre-norm with RMSNorm
+    - Bias-free linears
+    """
+    def __init__(self, hidden_size: int, intermediate_size: int, dropout: float = 0.0, dtype: torch.dtype = torch.float32):
+        super().__init__()
+        self.norm = nn.RMSNorm(hidden_size, eps=1e-6, dtype=dtype)
+        self.gate = nn.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype)
+        self.up   = nn.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype)
+        self.down = nn.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype)
+        self.act  = nn.SiLU()
+        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.norm(x)
+        h = self.act(self.gate(h)) * self.up(h)  # SwiGLU
+        h = self.down(h)
+        h = self.drop(h)
+        return x + h
+
+class StandardFFNLayer(nn.Module):
+    """
+    Pre-norm RMSNorm, classic MLP:
+      y = x + Dropout( W2( Act( W1( RMSNorm(x) ) ) ) )
+    - No SwiGLU: single hidden nonlinearity (GELU/ReLU/SiLU)
+    - Bias-free linears (common in modern LLM FFNs)
+    """
+    def __init__(
+        self,
+        hidden_size: int,
+        intermediate_size: int,
+        dropout: float = 0.0,
+        dtype: torch.dtype = torch.float32,
+        activation: str = "gelu",
+    ):
+        super().__init__()
+        self.norm = nn.RMSNorm(hidden_size, eps=1e-6, dtype=dtype)
+        self.w1   = nn.Linear(hidden_size, intermediate_size, bias=False, dtype=dtype)
+        self.w2   = nn.Linear(intermediate_size, hidden_size, bias=False, dtype=dtype)
+        self.drop = nn.Dropout(dropout) if dropout > 0 else nn.Identity()
+
+        act = activation.lower()
+        if act == "gelu":
+            self.act = nn.GELU()
+        elif act == "relu":
+            self.act = nn.ReLU()
+        elif act == "silu":
+            self.act = nn.SiLU()
+        else:
+            raise ValueError(f"Unsupported activation: {activation}")
+
+    def forward(self, x: Tensor) -> Tensor:
+        h = self.norm(x)
+        h = self.act(self.w1(h))
+        h = self.w2(h)
+        h = self.drop(h)
+        return x + h
+
+class RegularMLP(nn.Module):
+    """
+    Qwen3-style stacked MLP operating at a fixed hidden size.
+    - No input/output projections; caller is responsible for projections.
+    - num_layers repeats of Qwen-style FFN sublayer (pre-RMSNorm, SwiGLU, bias-free)
+    """
+    def __init__(
+        self,
+        hidden_dim: int = 1024,
+        intermediate_dim: int = 3072,
+        num_layers: int = 3,
+        dropout: float = 0.1,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+        assert num_layers >= 1, "num_layers must be >= 1"
+
+        self.blocks = nn.ModuleList([
+            StandardFFNLayer(hidden_size=hidden_dim, intermediate_size=intermediate_dim, dropout=dropout, dtype=dtype)
+            for _ in range(num_layers)
+        ])
+
+    def forward(self, x: Tensor) -> Tensor:
+        for blk in self.blocks:
+            x = blk(x)
+        return x
+
+@register_model
+@capture_init_args
+class C2CProjector(Projector):
+    """
+    Concise projector specialized to a fixed C2C configuration using StandardMLP.
+    - Projections: StandardMLP (pre-RMSNorm, SwiGLU, residual per sublayer)
+    - Concat: enabled, followed by linear combiner to target size
+    - Gate: scalar parameter with Gumbel-sigmoid during training
+    - Weights: input-dependent, head_merged granularity using target and projected key
+    - Target preservation: add_self=True, preserve_target_weight=False
+    - Temperatures: annealed gate temperature (1.0 -> 0.001 over 1929 steps), scalar_temperature=1.0
+    """
+
+    def __init__(
+        self,
+        source_dim: int,
+        target_dim: int,
+        source_num_heads: int = 1,
+        target_num_heads: int = 1,
+        intermediate_dim: int = 1024,
+        hidden_dim: int = 1024,
+        num_layers: int = 3,
+        dropout: float = 0.1,
+        initial_temperature: float = 1.0,
+        final_temperature: float = 0.001,
+        anneal_steps: int = 1929,
+        dtype: torch.dtype = torch.float32,
+    ):
+        super().__init__()
+
+        assert num_layers >= 3, "num_layers must be >= 3"
+
+        # Dimensions
+        self.source_dim = source_dim
+        self.target_dim = target_dim
+        self.source_num_heads = source_num_heads
+        self.target_num_heads = target_num_heads
+
+        # Sizes
+        in_dim = source_dim * source_num_heads
+        out_dim = target_dim * target_num_heads
+
+        # 1) concat(source_X, target_X) then project to hidden_dim
+        self.key_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype)
+        self.value_in = nn.Linear(in_dim + out_dim, hidden_dim, bias=True, dtype=dtype)
+
+        # 2) one-layer common embedding MLP to get intermediate representation (at hidden_dim)
+        self.key_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype)
+        self.value_mlp1 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=1, dropout=dropout, dtype=dtype)
+
+        # 3a) intermediate representation → (L-2)-layer MLP for weights → project to head dim
+        self.key_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype)
+        self.value_scalar_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=hidden_dim, num_layers=1, dropout=dropout, dtype=dtype)
+        self.key_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype)
+        self.value_scalar_head = nn.Linear(hidden_dim, target_num_heads, dtype=dtype)
+
+        # 3b) intermediate representation → (L-2)-layer MLP for projected_X → finally project hidden_dim → out_dim
+        self.key_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype)
+        self.value_proj_mlp2 = RegularMLP(hidden_dim=hidden_dim, intermediate_dim=intermediate_dim, num_layers=num_layers-2, dropout=dropout, dtype=dtype)
+        self.key_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype)
+        self.value_proj_out = nn.Linear(hidden_dim, out_dim, bias=True, dtype=dtype)
+
+        # Scalar key/value gate parameters and temperature schedule
+        self.key_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype))
+        self.value_gate_logit = nn.Parameter(torch.tensor(0.0, dtype=dtype))
+        self.use_gumbel = True
+        self.register_buffer("gate_temperature", torch.tensor(initial_temperature, dtype=dtype))
+        self.initial_temperature = initial_temperature
+        self.final_temperature = final_temperature
+        self.anneal_steps = anneal_steps
+
+        # Temperature for weight normalization
+        self.scalar_temperature = 1.0
+
+    def update_temperature(self, step: int):
+        ratio = min(step / self.anneal_steps, 1.0)
+        temp = self.initial_temperature * (self.final_temperature / self.initial_temperature) ** ratio
+        self.gate_temperature.fill_(temp)
+
+    def forward(
+        self,
+        source_kv: Tuple[Tensor, Tensor],
+        target_kv: Tuple[Tensor, Tensor],
+        position_ids: Optional[Tensor] = None,
+        max_pos: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        source_key, source_value = source_kv
+        target_key, target_value = target_kv
+
+        B, Hs, N, Ds = source_key.shape
+        _, Ht, _, Dt = target_key.shape
+
+        # Flatten heads
+        source_key_flat = source_key.transpose(1, 2).contiguous().view(B, N, Hs * Ds)
+        source_value_flat = source_value.transpose(1, 2).contiguous().view(B, N, Hs * Ds)
+        target_key_flat = target_key.transpose(1, 2).contiguous().view(B, N, Ht * Dt)
+        target_value_flat = target_value.transpose(1, 2).contiguous().view(B, N, Ht * Dt)
+
+        # 1) concat source and target features along channel
+        key_cat = torch.cat([source_key_flat, target_key_flat], dim=-1)
+        value_cat = torch.cat([source_value_flat, target_value_flat], dim=-1)
+
+        # 2) project to hidden dim
+        key_hidden = self.key_in(key_cat)
+        value_hidden = self.value_in(value_cat)
+
+        # 3) one-layer common embedding MLP to get intermediate representation (at hidden_dim)
+        key_hidden = self.key_mlp1(key_hidden)
+        value_hidden = self.value_mlp1(value_hidden)
+
+        # 4b) intermediate representation -> projected feature path
+        key_proj_hidden = self.key_proj_out(self.key_proj_mlp2(key_hidden)) # (B, N, Ht * Dt)
+        value_proj_hidden = self.value_proj_out(self.value_proj_mlp2(value_hidden)) # (B, N, Ht * Dt)
+        projected_key = key_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt)
+        projected_value = value_proj_hidden.view(B, N, Ht, Dt).transpose(1, 2) # (B, Ht, N, Dt)
+    
+        # 4a) intermediate representation -> scalar path
+        key_scalar = self.key_scalar_head(self.key_scalar_mlp2(key_hidden))       # (B, N, Ht)
+        value_scalar = self.value_scalar_head(self.value_scalar_mlp2(value_hidden)) # (B, N, Ht)
+        key_scalar = key_scalar.permute(0, 2, 1).unsqueeze(-1)   # (B, Ht, N, 1)
+        value_scalar = value_scalar.permute(0, 2, 1).unsqueeze(-1)  # (B, Ht, N, 1)
+
+        # Key/value gates: element-wise Gumbel noise with scalar logits (broadcast over channels)
+        key_gate_logit = self.key_gate_logit.view(1, 1, 1, 1)
+        value_gate_logit = self.value_gate_logit.view(1, 1, 1, 1)
+        if self.training and self.use_gumbel:
+            u1 = torch.rand(B, Ht, N, 1, device=key_gate_logit.device, dtype=key_gate_logit.dtype)
+            u2 = torch.rand(B, Ht, N, 1, device=value_gate_logit.device, dtype=value_gate_logit.dtype)
+            g1 = -torch.log(-torch.log(u1 + 1e-20) + 1e-20)
+            g2 = -torch.log(-torch.log(u2 + 1e-20) + 1e-20)
+            key_gate = torch.sigmoid((key_gate_logit + g1) / self.gate_temperature)
+            value_gate = torch.sigmoid((value_gate_logit + g2) / self.gate_temperature)
+        else:
+            key_gate = (key_gate_logit > 0).float()
+            value_gate = (value_gate_logit > 0).float()
+
+        # Normalize scalars (scalar_temperature=1.0)
+        norm_key_scalar = torch.sigmoid(key_scalar)
+        norm_value_scalar = torch.sigmoid(value_scalar)
+
+        # Combine (preserve_target_weight=False, add_self=True)
+        output_key = target_key + key_gate * norm_key_scalar * projected_key
+        output_value = target_value + value_gate * norm_value_scalar * projected_value
+
+        # Expose capture attributes for downstream analysis scripts
+        try:
+            # Store normalized scalars (detach to avoid autograd, keep device-agnostic via CPU)
+            self.last_norm_key_scalar = norm_key_scalar.detach().cpu()
+            self.last_norm_value_scalar = norm_value_scalar.detach().cpu()
+            # Store gate logits as python floats (parameters are scalar)
+            self.last_key_gate_logit = float(self.key_gate_logit.detach().cpu().item())
+            self.last_value_gate_logit = float(self.value_gate_logit.detach().cpu().item())
+        except Exception:
+            # Best-effort capture; never break forward path
+            pass
+
+        return output_key, output_value
+
+def save_projector(obj: Projector, file_path: str) -> None:
+    save_object(obj, file_path)
+
+def load_projector(file_path: str, override_args: Optional[dict] = None) -> Projector:
+    return load_object(file_path, get_projector_class, override_args)
+
+def create_projector(projector_type: str, **kwargs) -> Projector:
+    """
+    Factory function to create a projector based on type.
+    
+    Args:
+        projector_type: String indicating the type of projector
+        **kwargs: Additional arguments to pass to the projector constructor
+        
+    Returns:
+        An instance of the appropriate projector
+    """
+    # Prefer using the unified registry getter (handles case-insensitive keys)
+    try:
+        cls = get_projector_class(projector_type)
+    except ValueError as e:
+        raise e
+    return cls(**kwargs)

rosetta/model/sampling.py

@@ -0,0 +1,82 @@
+import torch
+import torch.nn.functional as F
+from typing import Union
+
+def sample_token(logits: torch.Tensor, temperature: float = 1.0, top_p: float = 1.0, top_k: int = -1) -> Union[int, torch.Tensor]:
+    """Sample a token from logits using temperature, top-p, and top-k sampling.
+    Args:
+        logits: Token logits of shape [vocab_size] or [batch_size, vocab_size]
+        temperature: Temperature for sampling (>0). Higher values produce more random samples.
+        top_p: Top-p probability threshold for nucleus sampling (0 < top_p ≤ 1)
+        top_k: Top-k threshold for sampling (if -1, no top-k filtering is applied)
+    Returns:
+        Sampled token ID (int for single sample, tensor for batch)
+    """
+    if not isinstance(logits, torch.Tensor):
+        raise TypeError("logits must be a torch.Tensor")
+    
+    if logits.dim() not in [1, 2]:
+        raise ValueError("logits must have shape [vocab_size] or [batch_size, vocab_size]")
+        
+    # Handle single dimension input
+    is_single_input = logits.dim() == 1
+    if is_single_input:
+        logits = logits.unsqueeze(0)
+    
+    batch_size = logits.shape[0]
+    
+    # For greedy sampling (temperature=0), just return argmax
+    if temperature == 0 or temperature <= 1e-5:
+        tokens = torch.argmax(logits, dim=-1)
+        return tokens.item() if is_single_input else tokens
+    
+    # Convert to probabilities
+    probs = torch.nn.functional.softmax(logits / temperature, dim=-1)
+    
+    # Apply top-k filtering first (if specified)
+    if top_k != -1:
+        # Get top-k values and indices
+        top_k_values, top_k_indices = torch.topk(probs, k=min(top_k, probs.shape[-1]), dim=-1)
+        
+        # Create a mask to zero out non-top-k probabilities
+        mask = torch.zeros_like(probs, dtype=torch.bool)
+        mask.scatter_(-1, top_k_indices, True)
+        
+        # Zero out non-top-k probabilities
+        probs = probs * mask.float()
+        
+        # Renormalize probabilities
+        probs = probs / probs.sum(dim=-1, keepdim=True)
+    
+    # Apply top-p (nucleus) sampling
+    if top_p < 1.0:
+        # Sort probabilities in descending order
+        sorted_probs, sorted_indices = torch.sort(probs, dim=-1, descending=True)
+        
+        # Calculate cumulative probabilities
+        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
+        
+        # Create a mask for probabilities to keep
+        # Values above top_p threshold are masked out
+        mask = cumulative_probs <= top_p
+        
+        # Always keep at least one token
+        mask[:, 0] = True
+        
+        # Zero out masked positions to exclude them from sampling
+        sorted_probs = sorted_probs * mask.float()
+        
+        # Renormalize probabilities
+        sorted_probs = sorted_probs / sorted_probs.sum(dim=-1, keepdim=True)
+        
+        # Sample from the filtered distribution
+        sampled_indices = torch.multinomial(sorted_probs, num_samples=1)
+        
+        # Map back to original vocabulary indices
+        tokens = torch.gather(sorted_indices, dim=-1, index=sampled_indices)
+        tokens = tokens.squeeze(-1)  # Remove sample dimension
+    else:
+        # Direct sampling if no top-p filtering
+        tokens = torch.multinomial(probs, num_samples=1).squeeze(-1)
+    
+    return tokens.item() if is_single_input else tokens

rosetta/model/wrapper.py

@@ -0,0 +1,840 @@
+"""
+The ensemble of multiple standard transformers LLM models, with automatic kv-cache projection. It shares the same interface as the standard transformers LLM models.
+"""
+
+from typing import List, Optional, Union
+import torch
+from torch import nn
+from transformers.cache_utils import Cache, DynamicCache
+from transformers.modeling_utils import PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+import json
+
+from rosetta.model.projector import Projector
+from rosetta.model.sampling import sample_token
+from transformers.utils import ModelOutput
+try:
+    from transformers.generation.utils import GreedySearchDecoderOnlyOutput, SampleDecoderOnlyOutput
+except Exception:
+    GreedySearchDecoderOnlyOutput = None
+    SampleDecoderOnlyOutput = None
+
+def clone_kv_cache(kv_cache: DynamicCache) -> DynamicCache:
+        new_cache = DynamicCache()
+        for k, v in zip(kv_cache.key_cache, kv_cache.value_cache):
+            new_cache.key_cache.append(k.clone().detach())
+            new_cache.value_cache.append(v.clone().detach())
+        return new_cache
+
+def hybrid_to_dynamic(hybrid_cache):
+    if hybrid_cache is None:
+        return None
+    if isinstance(hybrid_cache, DynamicCache):
+        return hybrid_cache
+
+    # 手动从 HybridCache 提取
+    if hasattr(hybrid_cache, "key_cache") and hasattr(hybrid_cache, "value_cache"):
+        keys = hybrid_cache.key_cache
+        values = hybrid_cache.value_cache
+        assert len(keys) == len(values), "key/value 层数不一致"
+
+        legacy_cache = [(k, v) for k, v in zip(keys, values)]
+        return DynamicCache.from_legacy_cache(legacy_cache)
+
+    raise TypeError(f"Unsupported cache type: {type(hybrid_cache)}")
+
+class RosettaModel(nn.Module):
+    """
+    Drop in replacement for the standard transformers LLM models, like Qwen3ForCausalLM
+    """
+    def __init__(self, model_list: List[PreTrainedModel], base_model_idx = 0, projector_list: List[Projector] = [], aggregator_list: List[nn.Module] = []):
+        super().__init__()
+        # model list: a list of model, model 0 by default is the base model
+        # projector list: a list of projector
+        # standard init with additional model list parameter
+        # kv-cache dict: key (source_model_idx, target_model_idx), value (Cache), assume only convert at prefill with one type of model
+        # projector dict: key (source_model_idx, target_model_idx) value dict(key (source_model_layer_idx, M_target value )
+
+        self.base_model_idx = base_model_idx
+        self.model_list = nn.ModuleList(model_list)
+
+        device = model_list[base_model_idx].device
+        dtype = model_list[base_model_idx].dtype
+        self.projector_list = nn.ModuleList(projector_list).to(device=device, dtype=dtype)
+        self.aggregator_list = nn.ModuleList(aggregator_list).to(device=device, dtype=dtype)
+
+        self.projector_dict = {}
+        self.aggregator_dict = {}
+        self.kv_cache_dict = {}
+        self._generation_hook_handlers = []
+
+    @property
+    def device(self):
+        return self.model_list[self.base_model_idx].device
+    
+    def to(self, device):
+        """
+        Move the RosettaModel and all underlying models and projectors to the specified device.
+        """
+        super().to(device)
+        for model in self.model_list:
+            model.to(device)
+        for projector in self.projector_list:
+            projector.to(device)
+        for aggregator in self.aggregator_list:
+            aggregator.to(device)
+        return self
+        
+    # set projector 
+    def set_projector_config(self, 
+                        source_model_idx: int, 
+                        source_model_layer_idx: int, 
+                        target_model_idx: int,
+                        target_model_layer_idx: int, 
+                        projector_idx: int):
+        """
+        Set the projector configuration
+        Args:
+            source_model_idx: int, the index of the source model
+            source_model_layer_idx: int, the index of the source model layer
+            target_model_idx: int, the index of the target model
+            target_model_layer_idx: int, the index of the target model layer
+            projector_idx: int, the index of the projector
+
+        The projector dict structure supports multiple projectors per target layer.
+        Structure:
+        {
+            target_model_idx: {
+                source_model_idx: {
+                    target_model_layer_idx: [(source_model_layer_idx, projector_idx), ...]
+                }
+            }
+        }
+        Repeated calls for the same (target, source, target_layer) append additional pairs.
+        """
+
+        if target_model_idx not in self.projector_dict.keys():
+            self.projector_dict[target_model_idx] = {}
+        if source_model_idx not in self.projector_dict[target_model_idx].keys():
+            self.projector_dict[target_model_idx][source_model_idx] = {}
+        # Accumulate list of (source_layer, projector_idx) for this target layer
+        layer_entry = self.projector_dict[target_model_idx][source_model_idx].get(target_model_layer_idx)
+        if layer_entry is None:
+            self.projector_dict[target_model_idx][source_model_idx][target_model_layer_idx] = [(source_model_layer_idx, projector_idx)]
+        else:
+            layer_entry.append((source_model_layer_idx, projector_idx))
+
+
+    def load_projector(self, projector_list):
+        self.projector_list: List[Projector] = projector_list
+
+    def load_aggregator(self, aggregator_list):
+        self.aggregator_list: List[nn.Module] = aggregator_list
+
+
+    def get_projector(self, 
+                        source_model_idx, 
+                        source_model_layer_idx, 
+                        target_model_idx,
+                        target_model_layer_idx):
+        pair_list = self.projector_dict[target_model_idx][source_model_idx][target_model_layer_idx]
+        if len(pair_list) == 0:
+            raise ValueError("No projector configured for the given target layer")
+        # Prefer exact source layer match
+        for src_layer, projector_id in pair_list:
+            if src_layer == source_model_layer_idx:
+                return self.projector_list[projector_id]
+        # Fallback: return the first projector
+        return self.projector_list[pair_list[0][1]]
+
+    def set_aggregator_idx(self,
+                           source_model_idx: int,
+                           target_model_idx: int,
+                           target_model_layer_idx: int,
+                           aggregator_idx: int):
+        if target_model_idx not in self.aggregator_dict:
+            self.aggregator_dict[target_model_idx] = {}
+        if source_model_idx not in self.aggregator_dict[target_model_idx]:
+            self.aggregator_dict[target_model_idx][source_model_idx] = {}
+        self.aggregator_dict[target_model_idx][source_model_idx][target_model_layer_idx] = aggregator_idx
+
+
+    @staticmethod
+    def load_json(file_name):
+        with open(file_name, "r") as f:
+            result = json.load(f)
+        return result
+    
+    @staticmethod
+    def _convert_dict_keys_to_ints(obj):
+        """
+        Recursively convert dictionary keys that look like integers back to int.
+        This reverses json.dump's coercion of dict keys to strings.
+        """
+        if isinstance(obj, dict):
+            new_obj = {}
+            for key, value in obj.items():
+                if isinstance(key, str) and key.lstrip('-').isdigit():
+                    new_key = int(key)
+                else:
+                    new_key = key
+                new_obj[new_key] = RosettaModel._convert_dict_keys_to_ints(value)
+            return new_obj
+        if isinstance(obj, list):
+            return [RosettaModel._convert_dict_keys_to_ints(v) for v in obj]
+        return obj
+    
+    
+    def save_projector_config(self, file_name):
+        with open(file_name, "w") as f:
+            json.dump(self.projector_dict, f)
+
+    
+    def load_projector_config(self, config_path):
+        if config_path.endswith(".json"):
+            loaded = RosettaModel.load_json(config_path)
+            self.projector_dict = RosettaModel._convert_dict_keys_to_ints(loaded)
+
+    def save_aggregator_config(self, file_name):
+        with open(file_name, "w") as f:
+            json.dump(self.aggregator_dict, f)
+
+    def load_aggregator_config(self, config_path):
+        if config_path.endswith(".json"):
+            loaded = RosettaModel.load_json(config_path)
+            self.aggregator_dict = RosettaModel._convert_dict_keys_to_ints(loaded)
+
+
+    def set_kv_cache_dict(self, source_model_idx, target_model_idx, cache):
+        if target_model_idx not in self.kv_cache_dict.keys():
+            self.kv_cache_dict[target_model_idx] = {}
+        if cache is None:
+            # Initialize with a DynamicCache instead of RosettaCache for now
+            self.kv_cache_dict[target_model_idx][source_model_idx] = DynamicCache() # noqa, maybe we should use RosettaCache here
+        else:
+            self.kv_cache_dict[target_model_idx][source_model_idx] = cache
+
+    def forward(
+        self,
+        kv_cache_index: Optional[List] = None,
+        input_ids: Optional[Union[torch.LongTensor, List[torch.LongTensor]]] = None,
+        attention_mask: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        # **kwargs: Unpack[KwargsForCausalLM],
+        *args,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """
+        Forward pass
+        KVCache index is a list of tensors with shape (B, sec_seq_len, 2), indicating the source and target kv cache model index
+
+        If input_ids is LongTensor, default to same input ids for different models
+        If input_ids is Tuple, default to different input ids for different models.
+
+        No Rosetta: (-1, 0)
+        """
+        
+        # noqa
+        self.kv_cache_dict = dict()
+
+        # Handle different input formats: if input_ids is a list, use per-model inputs
+        if isinstance(input_ids, list):
+            # Use list format: different input_ids and attention_mask for each model
+            base_input_ids = input_ids[self.base_model_idx] if input_ids is not None else None
+            base_attention_mask = attention_mask[self.base_model_idx] if attention_mask is not None else None
+            _, seqlen = base_input_ids.size() if base_input_ids is not None else (0, 0)
+        else:
+            # Use tensor format: same input_ids and attention_mask for all models (backward compatibility)
+            base_input_ids = input_ids
+            base_attention_mask = attention_mask
+            _, seqlen = input_ids.size() if input_ids is not None else (0, 0)
+
+        num_sections = len(kv_cache_index) if kv_cache_index is not None else 1
+
+        section_lengths = [kv_cache_index[i].shape[1] for i in range(num_sections)] if kv_cache_index is not None else [seqlen]
+        section_starts = [0]
+        for l in section_lengths:
+            section_starts.append(section_starts[-1] + l)
+        
+        curr_base_kv_cache = past_key_values
+
+        if seqlen >= 1:
+            for i in range(num_sections):
+                start = section_starts[i]
+                end = section_starts[i + 1]
+                prefill_input_ids = base_input_ids[:, start:end] if base_input_ids is not None else None
+                prefill_attention_mask = base_attention_mask[:, :end] if base_attention_mask is not None else None
+                prefill_position_ids = position_ids[:, start:end] if position_ids is not None else None
+                prefill_labels = labels[:, start:end] if labels is not None else None
+
+                output = self.model_list[self.base_model_idx].forward(
+                    input_ids=prefill_input_ids,
+                    attention_mask=prefill_attention_mask, 
+                    position_ids=prefill_position_ids,
+                    past_key_values=curr_base_kv_cache,
+                    labels=prefill_labels,
+                    use_cache=use_cache, 
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    *args,
+                    **kwargs
+                )
+
+                if self.base_model_idx not in self.kv_cache_dict:
+                    self.kv_cache_dict[self.base_model_idx] = {}
+                if self.base_model_idx not in self.kv_cache_dict[self.base_model_idx]:
+                    self.kv_cache_dict[self.base_model_idx][self.base_model_idx] = None
+                self.kv_cache_dict[self.base_model_idx][self.base_model_idx] = output.past_key_values
+
+                curr_base_kv_cache: DynamicCache = output.past_key_values
+
+                if i != num_sections - 1:
+                    for source_model_idx in range(1, len(self.model_list)):
+                        if self.base_model_idx not in self.kv_cache_dict:
+                            self.kv_cache_dict[self.base_model_idx] = {}
+                        if source_model_idx not in self.kv_cache_dict[self.base_model_idx]:
+                            self.kv_cache_dict[self.base_model_idx][source_model_idx] = None
+
+                        # Get model-specific input_ids and attention_mask
+                        if isinstance(input_ids, list):
+                            source_input_ids = input_ids[source_model_idx]
+                            source_attention_mask = attention_mask[source_model_idx] if attention_mask is not None else None
+                            source_prefill_input_ids = source_input_ids[:, start:end] if source_input_ids is not None else None
+                            source_prefill_attention_mask = source_attention_mask[:, :end] if source_attention_mask is not None else None
+                        else:
+                            # Backward compatibility: use same input for all models
+                            source_prefill_input_ids = prefill_input_ids
+                            source_prefill_attention_mask = prefill_attention_mask
+
+                        model = self.model_list[source_model_idx]
+                        was_training = model.training
+                        had_gc = getattr(model, "is_gradient_checkpointing", False)
+
+                        try:
+                            if was_training:
+                                model.eval()
+                            if had_gc:
+                                model.gradient_checkpointing_disable()
+
+                            with torch.no_grad():
+                                out = model(
+                                    input_ids=source_prefill_input_ids,
+                                    attention_mask=source_prefill_attention_mask,
+                                    position_ids=prefill_position_ids,
+                                    past_key_values=self.kv_cache_dict[self.base_model_idx][source_model_idx],
+                                    use_cache=True,
+                                    return_dict=True,
+                                )
+                                curr_source_kv_cache = out.past_key_values
+                        finally:
+                            if had_gc:
+                                model.gradient_checkpointing_enable()
+                            if was_training:
+                                model.train()
+                        
+                        curr_source_kv_cache = hybrid_to_dynamic(curr_source_kv_cache)
+                        self.kv_cache_dict[self.base_model_idx][source_model_idx] = curr_source_kv_cache
+
+                # calculate source model kvcache and apply projections
+                if self.base_model_idx in self.projector_dict:
+                    source_model_idx = kv_cache_index[i][0][0][0].item()  # Get the source model index from the kv_cache_index
+                    if source_model_idx != -1:
+                        for target_layer_idx, entry in self.projector_dict[self.base_model_idx][source_model_idx].items():
+                            base_key_cache, base_value_cache = curr_base_kv_cache[target_layer_idx]
+                            new_base_key_cache = base_key_cache[:, :, start:end, :]
+                            new_base_value_cache = base_value_cache[:, :, start:end, :]
+                            new_base_kv_cache = (new_base_key_cache, new_base_value_cache)
+
+                            pair_list = entry
+
+                            projected_kv_list = []
+                            source_kv_list = []
+                            for source_model_layer_idx, projector_idx in pair_list:
+                                source_key_cache, source_value_cache = self.kv_cache_dict[self.base_model_idx][source_model_idx][source_model_layer_idx]
+                                new_source_key_cache = source_key_cache[:, :, start:end, :]
+                                new_source_value_cache = source_value_cache[:, :, start:end, :]
+                                new_source_kv_cache = (new_source_key_cache, new_source_value_cache)
+                                projected_key, projected_value = self.projector_list[projector_idx].forward(
+                                    new_source_kv_cache, # tuple of (key, value), each of shape (B, N, H, D)
+                                    new_base_kv_cache
+                                )
+                                projected_kv_list.append((projected_key, projected_value))
+                                source_kv_list.append(new_source_kv_cache)
+
+                            # Aggregate (fallback to first projector if no aggregator is available)
+                            use_aggregator = (
+                                len(projected_kv_list) > 1 and
+                                len(self.aggregator_list) > 0 and
+                                self.base_model_idx in self.aggregator_dict and
+                                source_model_idx in self.aggregator_dict[self.base_model_idx] and
+                                target_layer_idx in self.aggregator_dict[self.base_model_idx][source_model_idx]
+                            )
+
+                            if use_aggregator:
+                                aggregator_idx = self.aggregator_dict[self.base_model_idx][source_model_idx][target_layer_idx]
+                                agg_key, agg_value = self.aggregator_list[aggregator_idx].forward(
+                                    source_kv_list,
+                                    new_base_kv_cache,
+                                    projected_kv_list
+                                )
+                            else:
+                                # Fallback to first projector result when no aggregator is available
+                                agg_key, agg_value = projected_kv_list[0]
+
+                            # Update cache with aggregated result
+                            curr_base_kv_cache.key_cache[target_layer_idx][:, :, start:end, :] = agg_key
+                            curr_base_kv_cache.value_cache[target_layer_idx][:, :, start:end, :] = agg_value
+                        
+                        output.past_key_values = curr_base_kv_cache
+                                                                             
+        # use base model for decode phase
+        else:
+            # Handle list input format for decode phase as well
+            decode_input_ids = input_ids[self.base_model_idx] if isinstance(input_ids, list) else input_ids
+            decode_attention_mask = attention_mask[self.base_model_idx] if isinstance(attention_mask, list) else attention_mask
+            
+            output = self.model_list[self.base_model_idx].forward(
+                input_ids=decode_input_ids,
+                attention_mask=decode_attention_mask,
+                position_ids=position_ids,
+                past_key_values=curr_base_kv_cache,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                cache_position=cache_position,
+                *args,
+                **kwargs
+            )
+        return output
+    
+    def oracle_forward(
+        self,
+        kv_cache_index: Optional[List] = None,
+        input_ids: Optional[Union[torch.LongTensor, List[torch.LongTensor]]] = None,
+        attention_mask: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Cache] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        labels: Optional[torch.LongTensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        cache_position: Optional[torch.LongTensor] = None,
+        logits_to_keep: Union[int, torch.Tensor] = 0,
+        # **kwargs: Unpack[KwargsForCausalLM],
+        *args,
+        **kwargs,
+    ) -> CausalLMOutputWithPast:
+        """
+        Forward pass
+        KVCache index is a list of tensors with shape (B, sec_seq_len, 2), indicating the source and target kv cache model index
+
+        If input_ids is LongTensor, default to same input ids for different models
+        If input_ids is Tuple, default to different input ids for different models.
+
+        No Rosetta: (-1, 0)
+        """
+        
+        # noqa
+        self.kv_cache_dict = dict()
+
+        # Handle different input formats: if input_ids is a list, use per-model inputs
+        if isinstance(input_ids, list):
+            # Use list format: different input_ids and attention_mask for each model
+            base_input_ids = input_ids[self.base_model_idx] if input_ids is not None else None
+            base_attention_mask = attention_mask[self.base_model_idx] if attention_mask is not None else None
+            _, seqlen = base_input_ids.size() if base_input_ids is not None else (0, 0)
+        else:
+            # Use tensor format: same input_ids and attention_mask for all models (backward compatibility)
+            base_input_ids = input_ids
+            base_attention_mask = attention_mask
+            _, seqlen = input_ids.size() if input_ids is not None else (0, 0)
+
+        num_sections = len(kv_cache_index) if kv_cache_index is not None else 1
+
+        section_lengths = [kv_cache_index[i].shape[1] for i in range(num_sections)] if kv_cache_index is not None else [seqlen]
+        section_starts = [0]
+        for l in section_lengths:
+            section_starts.append(section_starts[-1] + l)
+        
+        curr_base_kv_cache = past_key_values
+
+        loss = nn.MSELoss()
+        loss_output = 0
+        if seqlen > 1:
+            for i in range(num_sections):
+                start = section_starts[i]
+                end = section_starts[i + 1]
+                prefill_input_ids = base_input_ids[:, start:end] if base_input_ids is not None else None
+                prefill_attention_mask = base_attention_mask[:, :end] if base_attention_mask is not None else None
+                prefill_position_ids = position_ids[:, start:end] if position_ids is not None else None
+                prefill_labels = labels[:, start:end] if labels is not None else None
+
+                # calculate target model kvcache
+                output = self.model_list[self.base_model_idx].forward(
+                    input_ids=prefill_input_ids,
+                    attention_mask=prefill_attention_mask, 
+                    position_ids=prefill_position_ids,
+                    past_key_values=curr_base_kv_cache,
+                    labels=prefill_labels,
+                    use_cache=use_cache, 
+                    output_attentions=output_attentions,
+                    output_hidden_states=output_hidden_states,
+                    *args,
+                    **kwargs
+                )
+
+                if self.base_model_idx not in self.kv_cache_dict:
+                    self.kv_cache_dict[self.base_model_idx] = {}
+                if self.base_model_idx not in self.kv_cache_dict[self.base_model_idx]:
+                    self.kv_cache_dict[self.base_model_idx][self.base_model_idx] = None
+                self.kv_cache_dict[self.base_model_idx][self.base_model_idx] = output.past_key_values
+
+                curr_base_kv_cache: DynamicCache = output.past_key_values
+
+                if i != num_sections - 1:
+                    for source_model_idx in range(1, len(self.model_list)):
+                        if self.base_model_idx not in self.kv_cache_dict:
+                            self.kv_cache_dict[self.base_model_idx] = {}
+                        if source_model_idx not in self.kv_cache_dict[self.base_model_idx]:
+                            self.kv_cache_dict[self.base_model_idx][source_model_idx] = None
+
+                        # Get model-specific input_ids and attention_mask
+                        if isinstance(input_ids, list):
+                            source_input_ids = input_ids[source_model_idx]
+                            source_attention_mask = attention_mask[source_model_idx] if attention_mask is not None else None
+                            source_prefill_input_ids = source_input_ids[:, start:end] if source_input_ids is not None else None
+                            source_prefill_attention_mask = source_attention_mask[:, :end] if source_attention_mask is not None else None
+                        else:
+                            # Backward compatibility: use same input for all models
+                            source_prefill_input_ids = prefill_input_ids
+                            source_prefill_attention_mask = prefill_attention_mask
+
+                        curr_source_kv_cache = self.model_list[source_model_idx].forward(
+                            input_ids=source_prefill_input_ids,
+                            attention_mask=source_prefill_attention_mask,
+                            position_ids=prefill_position_ids,
+                            past_key_values=self.kv_cache_dict[self.base_model_idx][source_model_idx],
+                            use_cache=use_cache, 
+                            output_attentions=output_attentions,
+                            output_hidden_states=output_hidden_states,
+                            *args,
+                            **kwargs
+                        ).past_key_values
+                        self.kv_cache_dict[self.base_model_idx][source_model_idx] = curr_source_kv_cache
+
+                # calculate source model kvcache and apply projections
+                if self.base_model_idx in self.projector_dict:
+                    source_model_idx = kv_cache_index[i][0][0][0].item()  # Get the source model index from the kv_cache_index
+                    if source_model_idx != -1:
+                        for target_layer_idx, entry in self.projector_dict[self.base_model_idx][source_model_idx].items():
+                            base_key_cache, base_value_cache = curr_base_kv_cache[target_layer_idx]
+                            new_base_key_cache = base_key_cache[:, :, start:end, :]
+                            new_base_value_cache = base_value_cache[:, :, start:end, :]
+                            new_base_kv_cache = (new_base_key_cache, new_base_value_cache)
+
+                            pair_list = entry
+
+                            projected_kv_list = []
+                            source_kv_list = []
+                            for source_model_layer_idx, projector_idx in pair_list:
+                                source_key_cache, source_value_cache = self.kv_cache_dict[self.base_model_idx][source_model_idx][source_model_layer_idx]
+                                new_source_key_cache = source_key_cache[:, :, start:end, :]
+                                new_source_value_cache = source_value_cache[:, :, start:end, :]
+                                new_source_kv_cache = (new_source_key_cache, new_source_value_cache)
+                                projected_key, projected_value = self.projector_list[projector_idx].forward(
+                                    new_source_kv_cache, # tuple of (key, value), each of shape (B, N, H, D)
+                                    new_base_kv_cache
+                                )
+                                loss_output = loss_output + loss(torch.dstack([projected_key, projected_value]), 
+                                                   torch.dstack([new_source_key_cache, new_source_value_cache]))
+                                projected_kv_list.append((projected_key, projected_value))
+                                source_kv_list.append(new_source_kv_cache)
+
+                            # Aggregate (fallback to first projector if no aggregator is available)
+                            use_aggregator = (
+                                len(projected_kv_list) > 1 and
+                                len(self.aggregator_list) > 0 and
+                                self.base_model_idx in self.aggregator_dict and
+                                source_model_idx in self.aggregator_dict[self.base_model_idx] and
+                                target_layer_idx in self.aggregator_dict[self.base_model_idx][source_model_idx]
+                            )
+
+                            if use_aggregator:
+                                aggregator_idx = self.aggregator_dict[self.base_model_idx][source_model_idx][target_layer_idx]
+                                agg_key, agg_value = self.aggregator_list[aggregator_idx].forward(
+                                    source_kv_list,
+                                    new_base_kv_cache,
+                                    projected_kv_list
+                                )
+                            else:
+                                # Fallback to first projector result when no aggregator is available
+                                agg_key, agg_value = projected_kv_list[0]
+
+                            # Update cache with aggregated result
+                            curr_base_kv_cache.key_cache[target_layer_idx][:, :, start:end, :] = agg_key
+                            curr_base_kv_cache.value_cache[target_layer_idx][:, :, start:end, :] = agg_value
+                        
+                        output.past_key_values = curr_base_kv_cache
+                                                                             
+        # use base model for decode phase
+        else:
+            # Handle list input format for decode phase as well
+            decode_input_ids = input_ids[self.base_model_idx] if isinstance(input_ids, list) else input_ids
+            decode_attention_mask = attention_mask[self.base_model_idx] if isinstance(attention_mask, list) else attention_mask
+            
+            output = self.model_list[self.base_model_idx].forward(
+                input_ids=decode_input_ids,
+                attention_mask=decode_attention_mask,
+                position_ids=position_ids,
+                past_key_values=curr_base_kv_cache,
+                inputs_embeds=inputs_embeds,
+                labels=labels,
+                use_cache=use_cache,
+                output_attentions=output_attentions,
+                output_hidden_states=output_hidden_states,
+                cache_position=cache_position,
+                *args,
+                **kwargs
+            )
+
+        return output, loss_output
+    
+    @torch.no_grad()
+    def generate(
+        self,
+        kv_cache_index,
+        input_ids,
+        max_new_tokens: Optional[int] = None,
+        past_key_values: Optional[Cache] = None,
+        attention_mask: Optional[Union[torch.Tensor, List[torch.Tensor]]] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        eos_token_id: Optional[Union[int, List[int]]] = None,
+        pad_token_id: Optional[int] = None,
+        temperature: float = 1.0,
+        top_p: float = 1.0,
+        top_k: int = -1,
+        repetition_penalty: float = 1.0,
+        presence_penalty: float = 0.0,
+        frequency_penalty: float = 0.0,
+        do_sample: Optional[bool] = None,
+        return_dict_in_generate: Optional[bool] = None,
+        output_scores: Optional[bool] = None,
+        max_length: Optional[int] = None,
+        use_cache: bool = True,
+        streamer = None,
+        *args,
+        **kwargs,
+    ):
+        """
+        New generation loop without using the base model's generate.
+        - Uses this module's forward for prefill and per-token decode.
+        - Samples tokens via rosetta.model.sampling.sample_token.
+        Returns a tensor of shape [batch, prompt_len + generated_len] for the base model stream.
+        """
+        # Derive number of tokens to generate
+        # If max_new_tokens not provided, infer from max_length
+        if isinstance(input_ids, list):
+            base_input_ids_for_len = input_ids[self.base_model_idx]
+        else:
+            base_input_ids_for_len = input_ids
+        prompt_len = base_input_ids_for_len.size(1)
+
+        # Default eos/pad from base model tokenizer/config if not provided
+        base_model = self.model_list[self.base_model_idx]
+        gen_cfg = getattr(base_model, "generation_config", None)
+        cfg_obj = gen_cfg if gen_cfg is not None else getattr(base_model, "config", None)
+        if eos_token_id is None and cfg_obj is not None:
+            eos_token_id = getattr(cfg_obj, "eos_token_id", None)
+        if pad_token_id is None and cfg_obj is not None:
+            pad_token_id = getattr(cfg_obj, "pad_token_id", None)
+        if pad_token_id is None and eos_token_id is not None:
+            pad_token_id = eos_token_id if isinstance(eos_token_id, int) else eos_token_id[0]
+
+        if max_new_tokens is None:
+            if max_length is not None:
+                if max_length <= prompt_len:
+                    max_new_tokens = 0
+                else:
+                    max_new_tokens = max_length - prompt_len
+            else:
+                raise ValueError("Provide max_new_tokens or max_length")
+        if max_new_tokens < 0:
+            raise ValueError("max_new_tokens must be non-negative")
+
+        # Resolve base inputs
+        if isinstance(input_ids, list):
+            base_input_ids = input_ids[self.base_model_idx]
+            base_attention_mask = attention_mask[self.base_model_idx] if attention_mask is not None else None
+        else:
+            base_input_ids = input_ids
+            base_attention_mask = attention_mask
+
+        if base_attention_mask is None:
+            base_attention_mask = torch.ones_like(base_input_ids, dtype=torch.long, device=base_input_ids.device)
+
+        batch_size = base_input_ids.size(0)
+
+        # Prefill to build caches and obtain initial logits
+        prefill_output = self.forward(
+            kv_cache_index=kv_cache_index,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            *args,
+            **kwargs,
+        )
+
+        current_past = prefill_output.past_key_values
+        all_input_ids = base_input_ids
+        current_attention_mask = base_attention_mask
+
+        # Initialize streamer with prompt if provided
+        if streamer is not None:
+            streamer.put(base_input_ids)
+
+        # EOS handling setup
+        eos_set = None
+        if eos_token_id is not None:
+            eos_set = set(eos_token_id if isinstance(eos_token_id, list) else [eos_token_id])
+        finished = torch.zeros(batch_size, dtype=torch.bool, device=all_input_ids.device)
+
+        # Start from last prefill logits
+        last_logits = prefill_output.logits[:, -1, :]
+
+        # Determine sampling mode
+        if do_sample is None:
+            do_sample = False
+        effective_temperature = temperature if do_sample else 0.0
+
+        # Optional scores collection
+        collect_scores = bool(return_dict_in_generate) and bool(output_scores)
+        scores = []
+
+        for _ in range(max_new_tokens):
+            if collect_scores:
+                scores.append(last_logits)
+            # Apply repetition/presence/frequency penalties to logits before sampling
+            adjusted_logits = last_logits
+            if (
+                (repetition_penalty is not None and repetition_penalty != 1.0) or
+                (presence_penalty is not None and presence_penalty != 0.0) or
+                (frequency_penalty is not None and frequency_penalty != 0.0)
+            ):
+                adjusted_logits = last_logits.clone()
+                vocab_size = adjusted_logits.size(-1)
+                # Per-batch penalty application for clarity and correctness
+                for b in range(batch_size):
+                    seq_tokens = all_input_ids[b]
+                    if seq_tokens.numel() == 0:
+                        continue
+                    counts = torch.bincount(seq_tokens, minlength=vocab_size)
+                    if counts.dtype != torch.float32 and counts.dtype != torch.float64:
+                        counts = counts.to(adjusted_logits.dtype)
+                    # Presence penalty: penalize any token that has appeared
+                    if presence_penalty and presence_penalty != 0.0:
+                        presence_mask = counts > 0
+                        if presence_mask.any():
+                            adjusted_logits[b, presence_mask] = adjusted_logits[b, presence_mask] - presence_penalty
+                    # Frequency penalty: penalize proportionally to frequency
+                    if frequency_penalty and frequency_penalty != 0.0:
+                        adjusted_logits[b] = adjusted_logits[b] - frequency_penalty * counts
+                    # Repetition penalty (HF-style): divide positive logits, multiply negative logits
+                    if repetition_penalty and repetition_penalty != 1.0:
+                        rep_mask = counts > 0
+                        if rep_mask.any():
+                            pos_mask = rep_mask & (adjusted_logits[b] > 0)
+                            neg_mask = rep_mask & ~pos_mask
+                            if pos_mask.any():
+                                adjusted_logits[b, pos_mask] = adjusted_logits[b, pos_mask] / repetition_penalty
+                            if neg_mask.any():
+                                adjusted_logits[b, neg_mask] = adjusted_logits[b, neg_mask] * repetition_penalty
+
+            # Sample next token
+            next_token = sample_token(adjusted_logits, temperature=effective_temperature, top_p=top_p, top_k=top_k)
+            if not isinstance(next_token, torch.Tensor):
+                next_token = torch.tensor([next_token], device=all_input_ids.device, dtype=torch.long).repeat(batch_size)
+
+            # Apply EOS logic
+            if eos_set is not None:
+                just_finished = torch.zeros_like(finished)
+                for eid in eos_set:
+                    just_finished |= (next_token == eid)
+                finished = finished | just_finished
+                if pad_token_id is not None:
+                    next_token = torch.where(
+                        finished,
+                        torch.tensor(pad_token_id, device=next_token.device, dtype=next_token.dtype),
+                        next_token,
+                    )
+
+            # Append sampled token
+            next_token_unsqueezed = next_token.unsqueeze(1)
+            all_input_ids = torch.cat([all_input_ids, next_token_unsqueezed], dim=1)
+            current_attention_mask = torch.cat(
+                [
+                    current_attention_mask,
+                    torch.ones((batch_size, 1), device=current_attention_mask.device, dtype=current_attention_mask.dtype),
+                ],
+                dim=1,
+            )
+
+            # Stream the new token if streamer provided
+            if streamer is not None:
+                streamer.put(next_token_unsqueezed)
+
+            # Early stop if all sequences finished
+            if eos_set is not None and torch.all(finished):
+                break
+
+            # Decode one step using cached states; pass base-stream tensors
+            kv_cache_index = [torch.tensor([-1, 0], dtype=torch.long).repeat(1, 1).unsqueeze(0).to(all_input_ids.device)]
+
+            decode_output = self.forward(
+                kv_cache_index=kv_cache_index,
+                input_ids=next_token_unsqueezed,
+                attention_mask=current_attention_mask,
+                position_ids=None,
+                past_key_values=current_past,
+                use_cache=True,
+                *args,
+                **kwargs,
+            )
+            current_past = decode_output.past_key_values
+            last_logits = decode_output.logits[:, -1, :]
+
+        # End streaming if streamer provided
+        if streamer is not None:
+            streamer.end()
+
+        # Return style compatible with HF generate
+        if return_dict_in_generate:
+            if GreedySearchDecoderOnlyOutput is not None and SampleDecoderOnlyOutput is not None:
+                if do_sample:
+                    return SampleDecoderOnlyOutput(
+                        sequences=all_input_ids,
+                        scores=scores if collect_scores else None,
+                    )
+                else:
+                    return GreedySearchDecoderOnlyOutput(
+                        sequences=all_input_ids,
+                        scores=scores if collect_scores else None,
+                    )
+            # Fallback to generic ModelOutput
+            result = {"sequences": all_input_ids}
+            if collect_scores:
+                result["scores"] = scores
+            return ModelOutput(**result)
+        return all_input_ids

rosetta/train/__init__.py

@@ -0,0 +1,20 @@
+"""
+Training utilities for RosettaModel
+"""
+
+from .dataset_adapters import (
+    ChatDataset, 
+    RosettaDataCollator,
+)
+from .model_utils import setup_models
+
+__all__ = [
+    "RosettaTrainer",
+    "ProjectorSaveCallback", 
+    "freeze_model_components",
+    "InstructCoderChatDataset",
+    "ChatDataset",
+    "RosettaDataCollator",
+    "create_instructcoder_dataset",
+    "setup_models"
+] 

rosetta/train/dataset_adapters.py

@@ -0,0 +1,1648 @@
+"""
+Simple dataset adapter for converting InstructCoder to chat format
+"""
+
+from typing import List, Dict, Any, Optional, Union, Callable
+from datasets import load_dataset, load_from_disk
+from torch.utils.data import Dataset
+import torch
+from transformers import AutoTokenizer
+import inspect
+import os
+import hashlib
+# Dataset Registry System
+DATASET_REGISTRY = {}
+
+def register_dataset(cls=None, name=None):
+    """
+    Register a dataset class in the global registry.
+    Can be used as a decorator with or without arguments.
+    
+    Args:
+        cls: The class to register
+        name: Optional name to register the class under. If None, uses the class name.
+        
+    Returns:
+        The registered class
+    """
+    def _register(cls):
+        dataset_name = name if name is not None else cls.__name__
+        DATASET_REGISTRY[dataset_name] = cls
+        # Also register with lowercase name for case-insensitive lookup
+        DATASET_REGISTRY[dataset_name.lower()] = cls
+        return cls
+    
+    # Called as @register_dataset
+    if cls is not None:
+        return _register(cls)
+    
+    # Called as @register_dataset() or @register_dataset(name="DatasetName")
+    return _register
+
+
+def capture_init_args(cls):
+    """
+    Decorator to capture initialization arguments of a dataset class.
+    
+    Args:
+        cls: The class to decorate
+        
+    Returns:
+        The decorated class with automatic init args capture
+    """
+    original_init = cls.__init__
+    
+    def new_init(self, *args, **kwargs):
+        # Store all initialization arguments
+        self._init_args = {}
+        
+        # Get parameter names from the original __init__ method
+        sig = inspect.signature(original_init)
+        param_names = list(sig.parameters.keys())[1:]  # Skip 'self'
+        
+        # Map positional args to parameter names
+        for i, arg in enumerate(args):
+            if i < len(param_names):
+                self._init_args[param_names[i]] = arg
+        
+        # Add keyword args
+        self._init_args.update(kwargs)
+        
+        # Call the original __init__
+        original_init(self, *args, **kwargs)
+    
+    cls.__init__ = new_init
+    return cls
+
+
+# Unified batch filtering functions
+
+
+def create_text_length_filter(
+    max_length: int,
+    text_extractor: Callable[[Dict[str, Any]], str],
+    tokenizer: Optional[Any] = None,
+    use_tokens: bool = False
+):
+    """
+    Unified text length filter that can handle both word count and token count filtering.
+    
+    Args:
+        max_length: Maximum allowed length (words or tokens)
+        text_extractor: Function that extracts text from a single sample
+        tokenizer: Tokenizer for token counting (required if use_tokens=True)
+        use_tokens: If True, count tokens; if False, count words
+        
+    Returns:
+        Filter function that can be used with dataset.filter(batched=True)
+    """
+    if use_tokens and tokenizer is None:
+        raise ValueError("Tokenizer must be provided when use_tokens=True")
+    
+    def _text_length_filter_batch(batch):
+        batch_size = len(next(iter(batch.values())))
+        samples = [{key: values[i] for key, values in batch.items()} for i in range(batch_size)]
+        try:
+            texts = [text_extractor(sample) for sample in samples]
+            if use_tokens:
+                if hasattr(tokenizer, 'apply_chat_template') and any(isinstance(t, list) for t in texts):
+                    rendered = []
+                    for t in texts:
+                        if isinstance(t, list):
+                            rendered.append(tokenizer.apply_chat_template(t, tokenize=False, add_generation_prompt=False))
+                        else:
+                            rendered.append(str(t))
+                    tokenized = tokenizer(rendered, add_special_tokens=False)
+                else:
+                    tokenized = tokenizer([str(t) for t in texts], add_special_tokens=False)
+                lengths = [len(ids) for ids in tokenized["input_ids"]]
+            else:
+                lengths = [len(str(t).split()) for t in texts]
+            return [length <= max_length for length in lengths]
+        except Exception as e:
+            print(f"Error in text length filter: {e}")
+            return [False] * batch_size
+    
+    return _text_length_filter_batch
+
+
+def create_field_value_filter(target_value: Any, field_name: str, comparison: str = 'equal'):
+    """
+    Unified field value filter for exact matching, language filtering, etc.
+    
+    Args:
+        target_value: Value to compare against
+        field_name: Field name to check
+        comparison: Type of comparison ('equal', 'not_equal', 'in', 'not_in')
+        
+    Returns:
+        Filter function that can be used with dataset.filter(batched=True)
+    """
+    def _field_value_filter_batch(batch):
+        field_values = batch.get(field_name, [])
+        
+        if comparison == 'equal':
+            return [value == target_value for value in field_values]
+        elif comparison == 'not_equal':
+            return [value != target_value for value in field_values]
+        elif comparison == 'in':
+            return [value in target_value for value in field_values]
+        elif comparison == 'not_in':
+            return [value not in target_value for value in field_values]
+        else:
+            raise ValueError(f"Unsupported comparison: {comparison}")
+    
+    return _field_value_filter_batch
+
+
+def create_modulo_filter(mod_base: int, exclude_values: Union[int, List[int]], field_name: str = '_id'):
+    """
+    Unified modulo filter for ID-based filtering.
+    
+    Args:
+        mod_base: Modulo base
+        exclude_values: Value(s) to exclude (can be single int or list)
+        field_name: Field name containing the ID
+        
+    Returns:
+        Filter function that can be used with dataset.filter(batched=True)
+    """
+    if isinstance(exclude_values, int):
+        exclude_values = [exclude_values]
+    
+    def _modulo_filter_batch(batch):
+        ids = batch.get(field_name, [])
+        results = []
+        
+        for _id in ids:
+            try:
+                # Try numeric conversion first
+                id_num = int(_id)
+                mod_result = id_num % mod_base
+            except (ValueError, TypeError):
+                # Use hash for non-numeric IDs
+                id_hash = hash(str(_id))
+                mod_result = id_hash % mod_base
+            
+            results.append(mod_result not in exclude_values)
+        
+        return results
+    
+    return _modulo_filter_batch
+
+
+def create_conversation_length_filter(min_messages: int, text_field: str = 'conversations'):
+    """
+    Unified conversation length filter for OpenHermes-style datasets.
+    
+    Args:
+        min_messages: Minimum number of messages required (excluding system messages)
+        text_field: Field name containing the conversation
+        
+    Returns:
+        Filter function that can be used with dataset.filter(batched=True)
+    """
+    def _conversation_length_filter_batch(batch):
+        conversations_list = batch.get(text_field, [])
+        results = []
+        
+        for conversations in conversations_list:
+            try:
+                # Extract messages (excluding system)
+                message_count = 0
+                for msg in conversations:
+                    role = msg.get('from') or msg.get('role')
+                    if role in ('human', 'user', 'gpt', 'assistant'):
+                        message_count += 1
+                
+                results.append(message_count > min_messages)
+            except Exception:
+                results.append(False)
+        
+        return results
+    
+    return _conversation_length_filter_batch
+
+
+# Text extraction functions for common dataset patterns
+def extract_mmlu_text(sample: Dict[str, Any], question_field: str = 'question', choices_field: str = 'choices') -> str:
+    """Extract text from MMLU-style samples"""
+    question = sample.get(question_field, '')
+    choices = sample.get(choices_field, [])
+    
+    # Handle both list and dict formats for choices
+    if isinstance(choices, dict):
+        choices_text = choices.get('text', [])
+    else:
+        choices_text = choices
+    
+    return (str(question) + " " + " ".join(map(str, choices_text))).strip()
+
+
+def extract_chat_text(sample: Dict[str, Any], input_field: str = 'input', 
+                     context_field: str = 'context', answers_field: str = 'answers') -> List[Dict[str, str]]:
+    """Extract chat messages from LongBench-style samples"""
+    input_text = str(sample.get(input_field, ''))
+    context = str(sample.get(context_field, ''))
+    answers = sample.get(answers_field, [])
+    
+    assistant_message = answers[0] if answers and len(answers) > 0 else "No answer provided"
+    
+    # Build complete chat format
+    if context:
+        human_message = f"Context: {context}\n\nInstruction: {input_text}"
+    else:
+        human_message = f"Instruction: {input_text}"
+    
+    return [
+        {"role": "user", "content": human_message.strip()},
+        {"role": "assistant", "content": assistant_message.strip()}
+    ]
+
+
+def extract_conversation_text(sample: Dict[str, Any], text_field: str = 'conversations') -> str:
+    """Extract text from OpenHermes-style conversation samples"""
+    conversations = sample.get(text_field, [])
+    
+    if conversations and len(conversations) > 0:
+        return conversations[0].get('value', '')
+    return ''
+
+
+def extract_first_user_message(sample: Dict[str, Any], text_field: str = 'conversations') -> str:
+    """Extract the first human/user message from conversation-style samples."""
+    conversations = sample.get(text_field, [])
+    for msg in conversations:
+        role = msg.get('from') or msg.get('role')
+        if role in ('human', 'user'):
+            return str(msg.get('value', ''))
+    # Fallback to first message if role tags are missing
+    if conversations:
+        return str(conversations[0].get('value', ''))
+    return ''
+
+
+def extract_first_assistant_message(sample: Dict[str, Any], text_field: str = 'conversations') -> str:
+    """Extract the first gpt/assistant message from conversation-style samples."""
+    conversations = sample.get(text_field, [])
+    for msg in conversations:
+        role = msg.get('from') or msg.get('role')
+        if role in ('gpt', 'assistant'):
+            return str(msg.get('value', ''))
+    # Fallback to second message if present
+    if len(conversations) > 1:
+        return str(conversations[1].get('value', ''))
+    return ''
+
+
+def extract_openhermes_messages(sample: Dict[str, Any], text_field: str = 'conversations') -> List[Dict[str, str]]:
+    """Build chat messages excluding system; include all human/user and gpt/assistant in order."""
+    conversation = sample.get(text_field, [])
+    messages: List[Dict[str, str]] = []
+    for msg in conversation:
+        role = msg.get('from') or msg.get('role')
+        if role == 'system':
+            continue
+        if role in ('human', 'user'):
+            messages.append({"role": "user", "content": str(msg.get('value', '')).strip()})
+        elif role in ('gpt', 'assistant'):
+            messages.append({"role": "assistant", "content": str(msg.get('value', ''))})
+    return messages
+
+
+def extract_instruction_text(sample: Dict[str, Any], instruction_field: str = 'instruction', 
+                           inputs_field: str = 'inputs') -> str:
+    """Extract text from Inkuba-style instruction samples"""
+    instruction = sample.get(instruction_field)
+    inputs = sample.get(inputs_field, '')
+    
+    if instruction is not None:
+        return str(instruction) + "\n\n" + str(inputs)
+    else:
+        return str(inputs)
+
+
+def extract_chat_pair_text(sample: Dict[str, Any], user_field: str = 'inputs', 
+                          assistant_field: str = 'targets') -> List[Dict[str, str]]:
+    """Extract chat messages from Aya-style samples"""
+    user_text = str(sample.get(user_field, ''))
+    assistant_text = str(sample.get(assistant_field, ''))
+    
+    return [
+        {"role": "user", "content": user_text.strip()},
+        {"role": "assistant", "content": assistant_text.strip()}
+    ]
+
+
+
+def extract_dolly_chat_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]:
+    """Extract chat messages from Dolly-style samples.
+
+    Fields:
+      - instruction: str
+      - context: str (may be empty)
+      - response: str
+      - category: optional, may be empty/missing
+    """
+    instruction = str(sample.get('instruction', '')).strip()
+    context = str(sample.get('context', '') or '').strip()
+    response = str(sample.get('response', '')).strip()
+
+    if context:
+        user_message = f"{context}\n\n{instruction}"
+    else:
+        user_message = f"{instruction}"
+
+    return [
+        {"role": "user", "content": user_message.strip()},
+        {"role": "assistant", "content": response}
+    ]
+
+
+def extract_mmmlu_chat_messages(sample: Dict[str, Any]) -> List[Dict[str, str]]:
+    """Extract chat messages from MMMLU-style samples (OpenAI/MMMLU)."""
+    choice_labels = ['A', 'B', 'C', 'D']
+
+    template = (
+            "Jibu kwa usahihi swali lifuatalo:\n\n"
+            "{{question}}\n\n"
+            "Chaguo:\n"
+            "{{choices}}\n\n"
+            "Maelekezo:\n"
+            "- Soma swali na chaguo zote kwa makini.\n"
+            "- Chagua jibu sahihi zaidi kati ya yaliyotolewa.\n"
+            "- Jibu TU kwa herufi (A, B, C, D) inayolingana na jibu sahihi.\n"
+            "- Usijumuishe maelezo, maandishi ya ziada, au alama yoyote ya uakifishaji.\n\n"
+            "Jibu lako:"
+        )
+
+    choices_text = ""
+    for label in choice_labels:
+        content = sample.get(label, '')
+        choices_text += f"{label}. {content}\n"
+
+    user_prompt = template.replace("{{choices}}", choices_text).replace("{{question}}", str(sample.get('Question', '')))
+
+    correct_label = sample.get('Answer', '')
+    correct_content = sample.get(correct_label, '')
+    assistant_response = f"**Jibu lako: {correct_label}. {correct_content}.**"
+
+    return [
+        {"role": "user", "content": user_prompt.strip()},
+        {"role": "assistant", "content": assistant_response}
+    ]
+
+
+
+
+def apply_batch_filters(dataset, filters: list, filter_descriptions: list = None, 
+                       batch_size: int = 4096, combine_filters: bool = True,
+                       num_proc: Optional[int] = None):
+    """
+    Apply multiple filters using native batched filtering for maximum performance.
+    
+    Args:
+        dataset: Dataset to filter
+        filters: List of batched filter functions
+        filter_descriptions: Optional list of descriptions for logging
+        batch_size: Batch size for filtering operations
+        combine_filters: If True, combine all filters into a single batched operation
+        
+    Returns:
+        Filtered dataset and original length
+    """
+    if not filters:
+        return dataset, len(dataset)
+    
+    original_len = len(dataset)
+    
+    if combine_filters and len(filters) > 1:
+        # Combine all filters into a single batched operation for maximum efficiency
+        def _combined_batch_filter(batch):
+            # Get results from all filters
+            filter_results = []
+            for filter_func in filters:
+                filter_results.append(filter_func(batch))
+            
+            # Combine results with AND logic
+            combined_results = []
+            batch_size = len(filter_results[0]) if filter_results else 0
+            
+            for i in range(batch_size):
+                combined_results.append(all(result[i] for result in filter_results))
+            
+            return combined_results
+        
+        # Apply combined filter in a single pass
+        filtered_dataset = dataset.filter(
+            _combined_batch_filter,
+            batched=True,
+            batch_size=batch_size,
+            num_proc=num_proc if num_proc and (num_proc or 0) > 1 else None,
+            desc="Combined batch filtering"
+        )
+        
+        # Print filtering results
+        final_len = len(filtered_dataset)
+        if original_len != final_len:
+            print(f"Applied combined batch filtering: {original_len} -> {final_len} samples")
+            if filter_descriptions:
+                for desc in filter_descriptions:
+                    print(f"  - {desc}")
+    
+    else:
+        # Apply each filter sequentially with batched processing
+        current_dataset = dataset
+        
+        for i, (filter_func, desc) in enumerate(zip(filters, filter_descriptions or [''] * len(filters))):
+            pre_filter_len = len(current_dataset)
+            
+            current_dataset = current_dataset.filter(
+                filter_func,
+                batched=True,
+                batch_size=batch_size,
+                num_proc=num_proc if num_proc and (num_proc or 0) > 1 else None,
+                desc=f"Filtering: {desc}" if desc else f"Filter {i+1}"
+            )
+            
+            post_filter_len = len(current_dataset)
+            if desc and pre_filter_len != post_filter_len:
+                print(f"  - {desc}: {pre_filter_len} -> {post_filter_len} samples")
+        
+        filtered_dataset = current_dataset
+        final_len = len(filtered_dataset)
+        if original_len != final_len:
+            print(f"Applied sequential batch filtering: {original_len} -> {final_len} samples")
+    
+    return filtered_dataset, original_len
+
+
+def generate_kv_cache_index(instruction_length: int, full_length: int) -> torch.tensor:
+    """
+    Generate KV cache index for the input sequence.
+    
+    Args:
+        instruction_length: Length of the instruction tokens
+        full_length: Total length of the full conversation tokens
+        
+    Returns:
+        Tensor with KV cache index
+    """
+    assert instruction_length <= full_length
+
+    instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(instruction_length - 1, 1)
+    label_index = torch.tensor([-1, 0], dtype=torch.long).repeat(full_length - instruction_length + 1, 1)
+
+    kv_cache_index = torch.cat([instruction_index, label_index], dim=0)  # shape: (seq_len, 2)
+
+    return kv_cache_index
+
+
+"""
+Instruction dataset
+
+Convert any form of inputs to standard message format
+"""
+
+@register_dataset
+@capture_init_args
+class LongBenchChatDataset(Dataset):
+    """LongBench数据集转换为LongBench原始格式"""
+    
+    def __init__(self, split: str = "test", num_samples: Optional[int] = None,
+                 dataset_name: Optional[str] = None, language: Optional[str] = None,
+                 max_word_count: Optional[int] = None, max_length: Optional[int] = 14000,
+                 use_longbench_e: bool = True, filter_mod4: bool = True):
+        """
+        初始化LongBench数据集
+        
+        Args:
+            split: 数据集分割 ("test" - LongBench主要使用test分割)
+            num_samples: 使用的样本数量 (None表示全部)
+            dataset_name: 特定数据集名称 (None表示所有数据集)
+            language: 语言过滤 ("en" 或 "zh")
+            max_word_count: 最大词数限制（用于英文文本）
+            max_length: 最大字符长度限制
+            use_longbench_e: 是否使用LongBench-E版本
+            filter_mod4: 是否过滤_id mod4余1的样本
+        """
+        print(f"Loading LongBench{' -E' if use_longbench_e else ''} dataset (split: {split}, dataset: {dataset_name})...")
+        
+        # LongBench包含的数据集列表
+        longbench_datasets = [
+            "narrativeqa", "qasper", "multifieldqa_en", "multifieldqa_zh", "hotpotqa", 
+            "2wikimqa", "musique", "dureader", "gov_report", "qmsum", "multi_news", 
+            "vcsum", "trec", "triviaqa", "samsum", "lsht", "passage_count", 
+            "passage_retrieval_en", "passage_retrieval_zh", "lcc", "repobench-p"
+        ]
+        
+        longbench_e_datasets = [
+            "qasper", "multifieldqa_en", "hotpotqa", "2wikimqa", "gov_report", 
+            "multi_news", "trec", "triviaqa", "samsum", "passage_count", 
+            "passage_retrieval_en", "lcc", "repobench-p"
+        ]
+        
+        target_datasets = longbench_e_datasets if use_longbench_e else longbench_datasets
+        
+        # 定义LongBench提示模板
+        self.dataset_prompt_formats = {
+    "narrativeqa": "You are given a story, which can be either a novel or a movie script, and a question. Answer the question asconcisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nStory: {context}\n\nNow, answer the question based on the story asconcisely as you can, using a single phrase if possible. Do not provide any explanation.\n\nQuestion: {input}\n\nAnswer:",
+    "qasper": "You are given a scientific article and a question. Answer the question as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write \"unanswerable\". If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". Do not provide any explanation.\n\nArticle: {context}\n\n Answer the question based on the above article as concisely as you can, using a single phrase or sentence if possible. If the question cannot be answered based on the information in the article, write \"unanswerable\". If the question is a yes/no question, answer \"yes\", \"no\", or \"unanswerable\". Do not provide any explanation.\n\nQuestion: {input}\n\nAnswer:",
+    "multifieldqa_en": "Read the following text and answer briefly.\n\n{context}\n\nNow, answer the following question based on the above text, only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "multifieldqa_zh": "阅读以下文字并用中文简短回答：\n\n{context}\n\n现在请基于上面的文章回答下面的问题，只告诉我答案，不要输出任何其他字词。\n\n问题：{input}\n回答：",
+    "hotpotqa": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nThe following are given passages.\n{context}\n\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "2wikimqa": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nThe following are given passages.\n{context}\n\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "musique": "Answer the question based on the given passages. Only give me the answer and do not output any other words.\n\nThe following are given passages.\n{context}\n\nAnswer the question based on the given passages. Only give me the answer and do not output any other words.\n\nQuestion: {input}\nAnswer:",
+    "dureader": "请基于给定的文章回答下述问题。\n\n文章：{context}\n\n请基于上述文章回答下面的问题。\n\n问题：{input}\n回答：",
+    "gov_report": "You are given a report by a government agency. Write a one-page summary of the report.\n\nReport:\n{context}\n\nNow, write a one-page summary of the report.\n\nSummary:",
+    "qmsum": "You are given a meeting transcript and a query containing a question or instruction. Answer the query in one or more sentences.\n\nTranscript:\n{context}\n\nNow, answer the query based on the above meeting transcript in one or more sentences.\n\nQuery: {input}\nAnswer:",
+    "multi_news": "You are given several news passages. Write a one-page summary of all news. \n\nNews:\n{context}\n\nNow, write a one-page summary of all the news.\n\nSummary:",
+    "vcsum": "下面有一段会议记录，请你阅读后，写一段总结，总结会议的内容。\n会议记录：\n{context}\n\n会议总结：",
+    "trec": "Please determine the type of the question below. Here are some examples of questions.\n\n{context}\n{input}",
+    "triviaqa": "Answer the question based on the given passage. Only give me the answer and do not output any other words. The following are some examples.\n\n{context}\n\n{input}",
+    "samsum": "Summarize the dialogue into a few short sentences. The following are some examples.\n\n{context}\n\n{input}",
+    "lsht": "请判断给定新闻的类别，下面是一些例子。\n\n{context}\n{input}",
+    "passage_count": "There are some paragraphs below sourced from Wikipedia. Some of them may be duplicates. Please carefully read these paragraphs and determine how many unique paragraphs there are after removing duplicates. In other words, how many non-repeating paragraphs are there in total?\n\n{context}\n\nPlease enter the final count of unique paragraphs after removing duplicates. The output format should only contain the number, such as 1, 2, 3, and so on.\n\nThe final answer is: ",
+    "passage_retrieval_en": "Here are 30 paragraphs from Wikipedia, along with an abstract. Please determine which paragraph the abstract is from.\n\n{context}\n\nThe following is an abstract.\n\n{input}\n\nPlease enter the number of the paragraph that the abstract is from. The answer format must be like \"Paragraph 1\", \"Paragraph 2\", etc.\n\nThe answer is: ",
+    "passage_retrieval_zh": "以下是若干段落文字，以及其中一个段落的摘要。请确定给定的摘要出自哪一段。\n\n{context}\n\n下面是一个摘要\n\n{input}\n\n请输入摘要所属段落的编号。答案格式必须是\"段落1\"，\"段落2\"等格式\n\n答案是：",
+    "lcc": "Please complete the code given below. \n{context}Next line of code:\n",
+    "repobench-p": "Please complete the code given below. \n{context}{input}Next line of code:\n"
+}
+        
+        # 定义不使用聊天模板的任务
+        #self.no_chat_template_tasks = ["trec", "triviaqa", "samsum", "lsht", "lcc", "repobench-p"]
+        self.no_chat_template_tasks=['']
+        self.use_longbench_e = use_longbench_e
+        self.max_length = max_length
+
+        if dataset_name:
+            if dataset_name not in target_datasets:
+                raise ValueError(f"Dataset {dataset_name} not found in LongBench{' -E' if use_longbench_e else ''}")
+            target_datasets = [dataset_name]
+            self.current_evaluating_subject = dataset_name
+        else:
+            self.current_evaluating_subject = None
+        
+        # 加载所有选定的数据集
+        all_data = []
+        for dataset in target_datasets:
+            try:
+                dataset_suffix = f"{dataset}_e" if use_longbench_e else dataset
+                data = load_dataset('THUDM/LongBench', dataset_suffix, split=split)
+                print(f"  Loaded {len(data)} samples from {dataset}")
+                
+                # 添加数据集名称标识
+                data = data.map(lambda x: {"dataset_source": dataset})
+                all_data.append(data)
+            except Exception as e:
+                print(f"Warning: Failed to load {dataset}: {e}")
+                continue
+        
+        if not all_data:
+            raise ValueError("No datasets were successfully loaded")
+        
+
+        from datasets import concatenate_datasets
+        self.dataset = concatenate_datasets(all_data)
+        
+
+
+
+
+        # mod4!=1
+        if filter_mod4:
+            original_len = len(self.dataset)
+            
+            def _mod4_not_1(example):
+                _id = example.get('_id', '')
+                id_hash = int(hashlib.sha256(str(_id).encode('utf-8')).hexdigest(), 16)
+                
+                return id_hash % 4 != 1
+            
+            self.dataset = self.dataset.filter(_mod4_not_1)
+            print(f"Filtered by _id mod4 != 1: {original_len} -> {len(self.dataset)} samples")
+        
+        # 限制样本数量
+        if num_samples and num_samples < len(self.dataset):
+            self.dataset = self.dataset.select(range(num_samples))
+            
+        print(f"Loaded total {len(self.dataset)} samples from LongBench{' -E' if use_longbench_e else ''}")    
+    def __len__(self):
+        return len(self.dataset)
+    
+    def _format_longbench_example(self, example: Dict[str, Any], tokenizer: AutoTokenizer) -> str:
+
+        # 1. 确定任务类型
+        dataset_source = example.get('dataset_source', '')
+        if self.current_evaluating_subject:
+            current_subject = self.current_evaluating_subject
+        else:
+            current_subject = dataset_source
+            
+        # 仅当字符串以"_e"结尾时才替换
+        import re
+        subject = re.sub(r"_e$", "", current_subject) if self.use_longbench_e else current_subject
+        
+        # 2. 获取提示模板
+        if subject not in self.dataset_prompt_formats:
+            subject = "narrativeqa"  # 默认模板
+        prompt_format = self.dataset_prompt_formats[subject]
+        
+        # 3. 直接使用**example展开所有字段
+        raw_prompt = prompt_format.format(**example)
+        
+        # 4. 超长截断逻辑
+        tokenized_raw = tokenizer(raw_prompt, truncation=False, return_tensors="pt").input_ids[0]
+        if len(tokenized_raw) > self.max_length:
+            half_len = int(self.max_length / 2)
+            raw_prompt = tokenizer.decode(tokenized_raw[:half_len], skip_special_tokens=True) + \
+                        tokenizer.decode(tokenized_raw[-half_len:], skip_special_tokens=True)
+        
+        # 5. 应用Chat Template
+
+        final_prompt = raw_prompt
+        print(len(tokenized_raw))
+        return final_prompt
+    
+    def __getitem__(self, idx):
+
+        sample = self.dataset[idx]
+        
+        # 格式化样本
+        formatted_prompt = self._format_longbench_example(sample, self.tokenizer)
+        
+        # 提取答案
+        answers = sample.get('answers', [])
+        assistant_message = answers[0] if answers and len(answers) > 0 else "No answer provided"
+        
+        return [
+            {
+                "role": "user",
+                "content": formatted_prompt.strip()
+            },
+            {
+                "role": "assistant", 
+                "content": assistant_message.strip()
+            }
+        ]
+
+@register_dataset
+@capture_init_args
+class MMLUChatDataset(Dataset):
+    """Simple MMLU dataset converted to chat format"""
+
+    def __init__(self, split: str = "train", num_samples: Optional[int] = None, max_word_count: Optional[int] = None):
+        """
+        Initialize the dataset
+
+        Args:
+            split: Dataset split
+            num_samples: Number of samples to use (None for all)
+            max_word_count: If set, drop samples whose question + all choices exceed this word count
+        """
+        print(f"Loading MMLU dataset (split: {split})...")
+        # Load dataset
+        dataset = load_dataset("cais/mmlu", "all")
+        dataset = dataset[split]
+
+        # Ensure we have a proper Dataset object
+        if hasattr(dataset, 'select'):
+            self.dataset = dataset
+        else:
+            raise ValueError(f"Unexpected dataset type: {type(dataset)}")
+
+        # Limit samples if specified
+        if num_samples and num_samples < len(self.dataset):
+            self.dataset = self.dataset.select(range(num_samples))
+            
+        # Apply total token length filtering on full chat (user + assistant)
+        if max_word_count is not None:
+            # Use a small tokenizer for speed; total token length = chat(user+assistant)
+            self._mmlu_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
+            extractor = lambda sample: self._build_chat_messages(sample)
+            filters = [create_text_length_filter(max_word_count, extractor, self._mmlu_tokenizer, use_tokens=True)]
+            filter_descriptions = [f"Token count filter (full chat): max {max_word_count}"]
+            self.dataset, _ = apply_batch_filters(self.dataset, filters, filter_descriptions)
+
+        print(f"Loaded {len(self.dataset)} samples")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+        return self._build_chat_messages(sample)
+
+    def _build_chat_messages(self, sample: Dict[str, Any]) -> List[Dict[str, str]]:
+        choice_labels = ['A', 'B', 'C', 'D']
+        question = sample.get('question', '')
+        choices_list = sample.get('choices', [])
+        user_prompt = f"Question: {question}\n\nChoices:\n"
+        for i, choice in enumerate(choices_list):
+            label = choice_labels[i] if i < len(choice_labels) else chr(65 + i)
+            user_prompt += f"{label}. {choice}\n"
+        ans_idx = sample.get('answer', 0)
+        if isinstance(ans_idx, str) and ans_idx.isdigit():
+            ans_idx = int(ans_idx)
+        ans_label = choice_labels[ans_idx] if 0 <= int(ans_idx) < len(choice_labels) else chr(65 + int(ans_idx))
+        assistant_text = f"The correct answer is {ans_label}."
+        return [
+            {"role": "user", "content": user_prompt.strip()},
+            {"role": "assistant", "content": assistant_text.strip()},
+        ]
+
+@register_dataset
+@capture_init_args
+class MMLUCotChatDataset(Dataset):
+    """Simple MMLUCot dataset converted to chat format"""
+
+    def __init__(self, split: str = "train", num_samples: Optional[int] = None):
+        """
+        Initialize the dataset
+
+        Args:
+            split: Dataset split
+            num_samples: Number of samples to use (None for all)
+        """
+        print(f"Loading MMLUCot dataset (split: {split})...")
+        # Load dataset
+        dataset = load_dataset("Brench/MMLU-Pro-CoT-Train-43K")
+        dataset = dataset[split]
+
+        # Ensure we have a proper Dataset object
+        if hasattr(dataset, 'select'):
+            self.dataset = dataset
+        else:
+            raise ValueError(f"Unexpected dataset type: {type(dataset)}")
+
+        # Limit samples if specified
+        if num_samples and num_samples < len(self.dataset):
+            self.dataset = self.dataset.select(range(num_samples))
+
+        print(f"Loaded {len(self.dataset)} samples")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+
+        user_prompt = sample['question'] + "\n"
+
+        assistant_response = sample['chain_of_thoughts']
+
+        return [
+            {
+                "role": "user",
+                "content": user_prompt.strip()
+            },
+            {
+                "role": "assistant",
+                "content": assistant_response
+            }
+        ]
+
+@register_dataset
+@capture_init_args
+class LLMGeneratedChatDataset(Dataset):
+    """Simple LLM Generated dataset converted to chat format"""
+
+    def __init__(self, split: str = "train", num_samples: Optional[int] = None, data_path: str = "./teacher_datasets/output/dataset_finished", max_word_count: Optional[int] = None):
+        """
+        Initialize the dataset
+
+        Args:
+            split: Dataset split
+            num_samples: Number of samples to use (None for all)
+        """
+        print(f"Loading LLMGeneratedCot dataset (split: {split})...")
+        # Load dataset
+        dataset = load_from_disk(data_path)
+
+        # Ensure we have a proper Dataset object
+        if hasattr(dataset, 'select'):
+            self.dataset = dataset
+        else:
+            raise ValueError(f"Unexpected dataset type: {type(dataset)}")
+
+        tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
+        
+        if max_word_count is not None:
+            original_len = len(self.dataset)
+            half = max_word_count // 2
+            def _under_token_limit(batch):
+                q = tokenizer(batch["input_text"], add_special_tokens=False, padding=False, truncation=False)
+                a = tokenizer(batch["model_response"], add_special_tokens=False, padding=False, truncation=False)
+                return [
+                    (len(q_ids) <= half) and (len(q_ids) + len(a_ids) <= max_word_count)
+                    for q_ids, a_ids in zip(q["input_ids"], a["input_ids"])
+                ]
+
+            self.dataset = self.dataset.filter(
+                _under_token_limit,
+                batched=True,
+                batch_size=2048,                    # 视显存/内存调大
+                num_proc=min(8, os.cpu_count() or 1),
+                load_from_cache_file=True,
+                desc=f"Filter max_word_count={max_word_count}",
+            )
+            print(f"Filtered by max_word_count={max_word_count}: {original_len} -> {len(self.dataset)} samples")
+
+        # Limit samples if specified
+        if num_samples and num_samples < len(self.dataset):
+            self.dataset = self.dataset.select(range(num_samples))
+
+        print(f"Loaded {len(self.dataset)} samples")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+
+        input_text = sample.get('input_text', '') or ''
+
+        # Parse question and choices from input_text, which is expected to contain a
+        # "Choices:" section followed by labeled options like "A. ..."
+        def _parse_question_and_choices(text: str):
+            lines = (text or '').splitlines()
+            # Find the line index for "Choices:" (case-insensitive, ignoring spaces)
+            choices_idx = -1
+            for i, line in enumerate(lines):
+                if line.strip().lower().startswith('choices'):
+                    choices_idx = i
+                    break
+
+            if choices_idx == -1:
+                # Fallback: no explicit Choices header found
+                question_part = text.strip()
+                return question_part, ''
+
+            question_part = '\n'.join(lines[:choices_idx]).strip()
+
+            # Collect labeled choices until blank line or instruction-like line
+            collected = []
+            for raw in lines[choices_idx + 1:]:
+                s = raw.strip()
+                if not s:
+                    # Stop on first blank after having collected at least one choice
+                    if collected:
+                        break
+                    else:
+                        continue
+                lower = s.lower()
+                # Stop when hitting instruction section common in prompts
+                if lower.startswith('instructions:') or lower.startswith("let's ") or lower.startswith('you must'):
+                    break
+                # Accept formats like "A. ..." or "A) ..."
+                if len(s) >= 3 and s[0] in 'ABCDEFGHIJ' and s[1] in ').' and s[2] == ' ':
+                    collected.append(s)
+                else:
+                    # If we've started collecting and this line doesn't look like a choice, stop
+                    if collected:
+                        break
+                    # Otherwise ignore preamble noise
+                    continue
+
+            choices_block = '\n'.join(collected).strip()
+            return question_part, choices_block
+
+        question, choices_block = _parse_question_and_choices(input_text)
+
+        # Rebuild user prompt using the evaluation CoT template
+        template = """Accurately answer the following question:
+
+{{question}}
+
+Choices:
+{{choices}}
+
+Instructions:
+- Carefully read the question and all options.
+- Let's think step by step and you must explain your reasoning briefly.
+- Then give the final answer.
+- Keep your response within 150 words."""
+
+        filled_prompt = (
+            template
+            .replace("{{question}}", question or '')
+            .replace("{{choices}}", choices_block or '')
+        )
+
+        user_prompt = filled_prompt.strip() + "\n"
+
+        assistant_response = sample['model_response']
+
+        return [
+            {
+                "role": "user",
+                "content": user_prompt.strip()
+            },
+            {
+                "role": "assistant",
+                "content": assistant_response
+            }
+        ]
+
+@register_dataset
+@capture_init_args
+class OpenBookChatDataset(Dataset):
+    """Simple OpenBook dataset converted to chat format"""
+
+    def __init__(self, split: str = "train", num_samples: Optional[int] = None):
+        """
+        Initialize the dataset
+
+        Args:
+            split: Dataset split
+            num_samples: Number of samples to use (None for all)
+        """
+        print(f"Loading OpenBook dataset (split: {split})...")
+        # Load dataset
+        dataset = load_dataset("allenai/openbookqa", "main")
+        dataset = dataset[split]
+
+        # Ensure we have a proper Dataset object
+        if hasattr(dataset, 'select'):
+            self.dataset = dataset
+        else:
+            raise ValueError(f"Unexpected dataset type: {type(dataset)}")
+
+        # Limit samples if specified
+        if num_samples and num_samples < len(self.dataset):
+            self.dataset = self.dataset.select(range(num_samples))
+
+        print(f"Loaded {len(self.dataset)} samples")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+        choice_labels = ['A', 'B', 'C', 'D']
+
+        user_prompt = (
+            f"Question: {sample['question_stem']}\n\n"
+            f"Choices:\n"
+        )
+        for idx, choice in enumerate(sample['choices']['text']):
+            label = choice_labels[idx]
+            user_prompt += f"{label}. {choice}\n"
+
+        correct_label = sample["answerKey"]
+        assistant_response = f"The correct answer is {correct_label}."
+
+        return [
+            {
+                "role": "user",
+                "content": user_prompt.strip()
+            },
+            {
+                "role": "assistant",
+                "content": assistant_response
+            }
+        ]
+
+@register_dataset
+@capture_init_args
+class OpenHermesChatDataset(Dataset):
+    """Simple general dataset converted to chat format"""
+
+    def __init__(self, split: str = "train", num_samples: Optional[int] = None, max_word_count: Optional[int] = None, min_conversation_turns: int = 0):
+        """
+        Initialize the dataset
+
+        Args:
+            split: Dataset split
+            num_samples: Number of samples to use (None for all)
+            max_word_count: Maximum token count for filtering
+            min_conversation_turns: Minimum number of conversation turns (default 3 for multi-turn conversations)
+        """
+        print(f"Loading OpenHermes dataset (split: {split})...")
+        # Load dataset
+        dataset = load_dataset("teknium/OpenHermes-2.5")
+        dataset = dataset[split]
+
+        # Ensure we have a proper Dataset object
+        if hasattr(dataset, 'select'):
+            self.dataset = dataset
+        else:
+            raise ValueError(f"Unexpected dataset type: {type(dataset)}")
+        
+        # Limit samples if specified
+        if num_samples and num_samples < len(self.dataset):
+            self.dataset = self.dataset.select(range(num_samples))
+
+        # Apply filters
+        filters = []
+        filter_descriptions = []
+
+        # Filter by minimum conversation length (exclude conversations with <= 2 messages)
+        if min_conversation_turns > 0:
+            filters.append(create_conversation_length_filter(min_conversation_turns - 1, 'conversations'))
+            filter_descriptions.append(f"Conversation length filter: min {min_conversation_turns} messages (multi-turn only)")
+
+        # Apply conversation-level token count filtering (all messages combined <= max_word_count)
+        if max_word_count is not None:
+            tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen3-0.6B")
+            extractor = lambda sample: extract_openhermes_messages(sample, 'conversations')
+            filters.append(create_text_length_filter(max_word_count, extractor, tokenizer, use_tokens=True))
+            filter_descriptions.append(f"Token count filter: max {max_word_count}")
+
+        # Apply all filters
+        if filters:
+            self.dataset, _ = apply_batch_filters(self.dataset, filters, filter_descriptions, num_proc=8)
+
+        print(f"Loaded {len(self.dataset)} samples")
+
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        sample = self.dataset[idx]
+        return extract_openhermes_messages(sample, 'conversations')
+    
+"""
+Chat dataset
+
+Convert standard message format to input_ids and labels
+"""
+class ChatDataset(Dataset):
+    """Dataset for chat format training with HuggingFace Trainer compatibility"""
+    
+    def __init__(self, chat_dataset, tokenizer: AutoTokenizer, max_length: int = 32768):
+        self.chat_dataset = chat_dataset
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+    
+    def __len__(self):
+        return len(self.chat_dataset)
+    
+    def __getitem__(self, idx) -> Dict[str, Any]:
+        messages = self.chat_dataset[idx]
+        
+        # Get instruction (first message)
+        instruction = self.tokenizer.apply_chat_template(
+            messages[:-1],
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+
+        # Get full conversation
+        full_text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+            enable_thinking=False,
+        )
+
+        # Tokenize instruction and full text
+        instruction_tokens = self.tokenizer(instruction, add_special_tokens=False)["input_ids"]
+        full_tokens = self.tokenizer(full_text, add_special_tokens=False)["input_ids"]
+        
+        # Truncate if necessary
+        if len(full_tokens) > self.max_length:
+            full_tokens = full_tokens[:self.max_length]
+        
+        # Create labels (-100 for instruction tokens, actual tokens for response)
+        labels = [-100] * len(instruction_tokens) + full_tokens[len(instruction_tokens):]
+        # labels = [-100] * (len(full_tokens) - 4) + full_tokens[-4:]
+        if len(labels) > self.max_length:
+            labels = labels[:self.max_length]
+        
+        kv_cache_index = generate_kv_cache_index(len(instruction_tokens), len(full_tokens))
+        # kv_cache_index = generate_kv_cache_index(len(full_tokens)-4, len(full_tokens))
+        # kv_cache_index = generate_kv_cache_index(len(full_tokens) + 1, len(full_tokens))
+
+        return {
+            "input_ids": full_tokens,
+            "labels": labels,
+            "kv_cache_index": kv_cache_index
+        }
+
+
+class AlignedChatDataset(Dataset):
+    """Dataset that precomputes aligned inputs for SLM/LLM using a TokenAligner"""
+    
+    def __init__(self, instruct_dataset: Dataset, aligner: Any, max_length: int = 32768):
+        self.dataset = instruct_dataset
+        self.aligner = aligner
+        self.max_length = max_length
+    
+    def __len__(self):
+        return len(self.dataset)
+    
+    def __getitem__(self, idx):
+        messages = self.dataset[idx]
+
+        # Build aligned sequences and section map
+        details = self.aligner.align_chat_messages(messages, add_generation_prompt=False, return_details=True)
+        slm_ids: List[int] = details['slm_ids_padded']
+        llm_ids: List[int] = details['llm_ids_padded']
+        sections = details['sections']
+
+        slm_pad_mask = torch.tensor(details['slm_padding_mask'])
+        llm_pad_mask = torch.tensor(details['llm_padding_mask'])
+        message_mask = torch.tensor(details['message_mask'])
+
+        # Determine instruction boundary as start of the last message section
+        instr_end = 0
+        for sec_idx in range(len(sections) - 1, -1, -1):
+            sec = sections[sec_idx]
+            if sec['type'] == 'message':
+                instr_end = sec['slm_range'][0]
+                break
+
+        # Labels: follow ChatDataset policy (-100 for instruction-only, supervise the rest)
+        labels = [-100] * instr_end + slm_ids[instr_end:]
+        if len(labels) > self.max_length:
+            labels = labels[:self.max_length]
+
+        # Truncate inputs if needed
+        if len(slm_ids) > self.max_length:
+            slm_ids = slm_ids[:self.max_length]
+            # Truncate padding mask accordingly
+            slm_pad_mask = slm_pad_mask[:self.max_length]
+        if len(llm_ids) > self.max_length:
+            llm_ids = llm_ids[:self.max_length]
+            llm_pad_mask = llm_pad_mask[:self.max_length]
+
+        # KV cache index based on instruction length
+        kv_cache_index = generate_kv_cache_index(instr_end, len(slm_ids))
+        # Addtionally mask non-message parts
+        kv_cache_index[~message_mask] = torch.tensor([[-1,0]])
+
+        return {
+            "input_ids": [slm_ids, llm_ids],
+            "labels": labels,
+            "kv_cache_index": kv_cache_index,
+            "messages": messages,
+            # Per-model aligned inputs (per-sample, pre-batch)
+            "model_padding_mask": [slm_pad_mask, llm_pad_mask],
+        }
+
+
+class BaselineChatDataset(Dataset):
+    """Simple dataset for baseline model training without Rosetta-specific features"""
+    
+    def __init__(self, chat_dataset, tokenizer: AutoTokenizer, max_length: int = 2048):
+        self.chat_dataset = chat_dataset
+        self.tokenizer = tokenizer
+        self.max_length = max_length
+    
+    def __len__(self):
+        return len(self.chat_dataset)
+    
+    def __getitem__(self, idx):
+        messages = self.chat_dataset[idx]
+        
+        # Get instruction (first message)
+        instruction = self.tokenizer.apply_chat_template(
+            messages[:1],
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False,
+        )
+
+        # Get full conversation
+        full_text = self.tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+            enable_thinking=False,
+        )
+
+        # Tokenize instruction and full text
+        instruction_tokens = self.tokenizer(instruction, add_special_tokens=False)["input_ids"]
+        full_tokens = self.tokenizer(full_text, add_special_tokens=False)["input_ids"]
+        
+        # Truncate if necessary
+        if len(full_tokens) > self.max_length:
+            full_tokens = full_tokens[:self.max_length]
+        
+        # Create labels (-100 for instruction tokens, actual tokens for response)
+        labels = [-100] * len(instruction_tokens) + full_tokens[len(instruction_tokens):]
+        if len(labels) > self.max_length:
+            labels = labels[:self.max_length]
+
+        return {
+            "input_ids": full_tokens,
+            "labels": labels,
+        }
+
+"""
+Data collator
+
+Batch chat data to model input
+"""
+
+class RosettaDataCollator:
+    """Improved data collator for RosettaModel training with cleaner logic"""
+
+    def __init__(self, slm_tokenizer: AutoTokenizer, llm_tokenizer: AutoTokenizer = None, 
+                 pad_to_multiple_of: Optional[int] = None, max_length: Optional[int] = None, 
+                 aligner: Optional[Any] = None, do_alignment: bool = False):
+        """
+        Initialize the collator.
+        
+        Args:
+            slm_tokenizer: Small language model tokenizer
+            llm_tokenizer: Large language model tokenizer (optional)
+            pad_to_multiple_of: Pad sequence length to multiple of this value
+            max_length: Maximum sequence length
+            aligner: Alignment module (if needed)
+            do_alignment: Whether to perform alignment
+        """
+        self.slm_tokenizer = slm_tokenizer
+        self.llm_tokenizer = llm_tokenizer
+        self.pad_to_multiple_of = pad_to_multiple_of
+        self.max_length = max_length
+        self.aligner = aligner
+        self.do_alignment = do_alignment
+        
+        if self.do_alignment:
+            assert self.aligner is not None, "Aligner must be provided if do_alignment is True"
+        
+        # Store padding token IDs for different models
+        self.slm_pad_token_id = self.slm_tokenizer.pad_token_id
+        self.llm_pad_token_id = self.llm_tokenizer.pad_token_id if self.llm_tokenizer else self.slm_pad_token_id
+
+    def _normalize_input_format(self, feature: Dict[str, Any]) -> Dict[str, Any]:
+        """
+        Normalize input format to handle both single and dual model inputs.
+        
+        Args:
+            feature: Input feature dictionary
+            
+        Returns:
+            Normalized feature with consistent format
+        """
+        # Normalize input_ids: ensure it's always a list of tensors
+        input_ids = feature['input_ids']
+        if isinstance(input_ids, list) and len(input_ids) > 0:
+            if isinstance(input_ids[0], list):
+                # Case: [[ids1], [ids2]] -> convert to list of tensors
+                input_ids_tensors = [torch.tensor(ids, dtype=torch.long) for ids in input_ids]
+            else:
+                # Case: [id1, id2, ...] -> single model case
+                input_ids_tensors = [torch.tensor(input_ids, dtype=torch.long)]
+        else:
+            # Fallback: assume single model
+            input_ids_tensors = [torch.tensor(input_ids, dtype=torch.long)]
+        
+        # Normalize attention_mask
+        attention_masks = []
+        if "model_padding_mask" in feature:
+            # Use model-specific padding masks
+            for model_padding_mask in feature["model_padding_mask"]:
+                attention_masks.append((~model_padding_mask).float())
+        else:
+            # Generate default attention masks
+            for input_tensor in input_ids_tensors:
+                attention_masks.append(torch.ones(len(input_tensor), dtype=torch.float))
+        
+        return {
+            'input_ids': input_ids_tensors,
+            'attention_mask': attention_masks,
+            'labels': torch.tensor(feature['labels'], dtype=torch.long),
+            'kv_cache_index': feature['kv_cache_index'],
+            'position_ids': torch.arange(len(feature['labels']), dtype=torch.long)
+        }
+
+    def _split_into_sections(self, normalized_feature: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+        Split sequence into sections based on kv_cache_index changes.
+        
+        Args:
+            normalized_feature: Normalized feature dictionary
+            
+        Returns:
+            List of sections
+        """
+        kv_idx = normalized_feature['kv_cache_index']
+        
+        # Find change points in kv_cache_index
+        change_points = [0]
+        for i in range(1, kv_idx.size(0)):
+            if not torch.equal(kv_idx[i], kv_idx[i - 1]):
+                change_points.append(i)
+        change_points.append(kv_idx.size(0))
+        
+        # Create sections
+        sections = []
+        for i in range(len(change_points) - 1):
+            start, end = change_points[i], change_points[i + 1]
+            section = {
+                'input_ids': [ids[start:end] for ids in normalized_feature['input_ids']],
+                'attention_mask': [mask[start:end] for mask in normalized_feature['attention_mask']],
+                'labels': normalized_feature['labels'][start:end],
+                'kv_cache_index': normalized_feature['kv_cache_index'][start:end],
+                'position_ids': normalized_feature['position_ids'][start:end]
+            }
+            sections.append(section)
+        
+        return sections
+
+    def _pad_sections(self, all_sections: List[List[Dict[str, Any]]]) -> Dict[str, Any]:
+        """
+        Pad sections to ensure uniform structure across batch.
+        
+        Args:
+            all_sections: List of section lists for each sample
+            
+        Returns:
+            Padded batch dictionary
+        """
+        max_sections = max(len(sections) for sections in all_sections)
+        num_models = len(all_sections[0][0]['input_ids']) if all_sections else 1
+        
+        # Initialize output structure - keep models separate throughout
+        padded_output = {
+            'input_ids_per_model': [[] for _ in range(num_models)],  # One list per model
+            'attention_mask_per_model': [[] for _ in range(num_models)],  # One list per model
+            'labels': [],
+            'kv_cache_index': [],
+            'position_ids': []
+        }
+        
+        # Process each section index
+        for sec_idx in range(max_sections):
+            section_data = self._collect_section_data(all_sections, sec_idx, num_models)
+            padded_section = self._pad_single_section(section_data, num_models)
+            
+            # Add to output - keep models separate
+            for model_idx in range(num_models):
+                padded_output['input_ids_per_model'][model_idx].append(
+                    padded_section['input_ids_per_model'][model_idx])
+                padded_output['attention_mask_per_model'][model_idx].append(
+                    padded_section['attention_mask_per_model'][model_idx])
+            
+            padded_output['labels'].append(padded_section['labels'])
+            padded_output['kv_cache_index'].append(padded_section['kv_cache_index'])
+            padded_output['position_ids'].append(padded_section['position_ids'])
+        
+        # Concatenate sections and finalize
+        return self._finalize_output(padded_output, num_models, len(all_sections))
+
+    def _collect_section_data(self, all_sections: List[List[Dict[str, Any]]], 
+                            sec_idx: int, num_models: int) -> Dict[str, List]:
+        """Collect data for a specific section across all samples."""
+        # Separate collections for each model to avoid confusion
+        section_data = {
+            'input_ids_per_model': [[] for _ in range(num_models)],  # [[slm_seqs], [llm_seqs]]
+            'attention_mask_per_model': [[] for _ in range(num_models)],
+            'labels': [],
+            'kv_cache_index': [],
+            'position_ids': []
+        }
+        
+        for sample_sections in all_sections:
+            # Some samples may have fewer sections; create default empty tensors when missing
+            if sec_idx < len(sample_sections):
+                sec = sample_sections[sec_idx]
+                for model_idx in range(num_models):
+                    section_data['input_ids_per_model'][model_idx].append(sec['input_ids'][model_idx])
+                    section_data['attention_mask_per_model'][model_idx].append(sec['attention_mask'][model_idx])
+                section_data['labels'].append(sec['labels'])
+                section_data['kv_cache_index'].append(sec['kv_cache_index'])
+                section_data['position_ids'].append(sec['position_ids'])
+            else:
+                # Default empty tensors; downstream pad_sequence will pad appropriately
+                for model_idx in range(num_models):
+                    section_data['input_ids_per_model'][model_idx].append(torch.tensor([], dtype=torch.long))
+                    section_data['attention_mask_per_model'][model_idx].append(torch.tensor([], dtype=torch.float))
+                section_data['labels'].append(torch.tensor([], dtype=torch.long))
+                section_data['kv_cache_index'].append(torch.empty((0, 2), dtype=torch.long))
+                section_data['position_ids'].append(torch.tensor([], dtype=torch.long))
+                
+        return section_data
+
+    def _pad_single_section(self, section_data: Dict[str, List], num_models: int) -> Dict[str, Any]:
+        """Pad tensors within a single section."""
+        # Pad input_ids separately for each model with their respective pad tokens
+        padded_input_ids_per_model = []
+        padded_attention_mask_per_model = []
+        
+        for model_idx in range(num_models):
+            pad_token_id = self.slm_pad_token_id if model_idx == 0 else self.llm_pad_token_id
+            
+            # Pad input_ids for this model
+            padded_input_ids = torch.nn.utils.rnn.pad_sequence(
+                section_data['input_ids_per_model'][model_idx], 
+                batch_first=True, 
+                padding_value=pad_token_id
+            )
+            padded_input_ids_per_model.append(padded_input_ids)
+            
+            # Pad attention_mask for this model
+            padded_attention_mask = torch.nn.utils.rnn.pad_sequence(
+                section_data['attention_mask_per_model'][model_idx],
+                batch_first=True,
+                padding_value=0
+            )
+            padded_attention_mask_per_model.append(padded_attention_mask)
+        
+        # Standard padding for other tensors
+        padded_labels = torch.nn.utils.rnn.pad_sequence(
+            section_data['labels'], batch_first=True, padding_value=-100)
+        padded_kv_cache = torch.nn.utils.rnn.pad_sequence(
+            section_data['kv_cache_index'], batch_first=True, padding_value=-1)
+        padded_position_ids = torch.nn.utils.rnn.pad_sequence(
+            section_data['position_ids'], batch_first=True, padding_value=0)
+        
+        return {
+            'input_ids_per_model': padded_input_ids_per_model,  # Keep separate per model
+            'attention_mask_per_model': padded_attention_mask_per_model,  # Keep separate per model
+            'labels': padded_labels,
+            'kv_cache_index': padded_kv_cache,
+            'position_ids': padded_position_ids,
+            'num_models': num_models
+        }
+
+    def _finalize_output(self, padded_output: Dict[str, List], 
+                        num_models: int, batch_size: int) -> Dict[str, Any]:
+        """Finalize the output by concatenating sections - keep models separate throughout."""
+        final_output = {}
+        
+        # Handle input_ids and attention_mask - keep separate per model
+        if num_models == 1:
+            # Single model case: concatenate sections for the single model
+            final_output['input_ids'] = torch.cat(padded_output['input_ids_per_model'][0], dim=1)
+            final_output['attention_mask'] = torch.cat(padded_output['attention_mask_per_model'][0], dim=1)
+        else:
+            # Multi-model case: keep as list of tensors, one per model
+            final_output['input_ids'] = [
+                torch.cat(padded_output['input_ids_per_model'][model_idx], dim=1) 
+                for model_idx in range(num_models)
+            ]
+            final_output['attention_mask'] = [
+                torch.cat(padded_output['attention_mask_per_model'][model_idx], dim=1)
+                for model_idx in range(num_models)
+            ]
+        
+        # Concatenate other tensors normally
+        final_output['labels'] = torch.cat(padded_output['labels'], dim=1)
+        final_output['position_ids'] = torch.cat(padded_output['position_ids'], dim=1)
+        final_output['kv_cache_index'] = padded_output['kv_cache_index']  # Keep as list of sections
+        
+        return final_output
+
+    def _apply_length_constraints(self, output: Dict[str, Any]) -> Dict[str, Any]:
+        """Apply max_length truncation if specified."""
+        if self.max_length is None:
+            return output
+        
+        # Determine current sequence length
+        if isinstance(output['input_ids'], list):
+            seq_length = output['input_ids'][0].size(1)
+        else:
+            seq_length = output['input_ids'].size(1)
+        
+        if seq_length <= self.max_length:
+            return output
+        
+        # Truncate sequences
+        if isinstance(output['input_ids'], list):
+            output['input_ids'] = [ids[:, :self.max_length] for ids in output['input_ids']]
+            output['attention_mask'] = [mask[:, :self.max_length] for mask in output['attention_mask']]
+        else:
+            output['input_ids'] = output['input_ids'][:, :self.max_length]
+            output['attention_mask'] = output['attention_mask'][:, :self.max_length]
+        
+        output['labels'] = output['labels'][:, :self.max_length]
+        output['position_ids'] = output['position_ids'][:, :self.max_length]
+        
+        # Truncate kv_cache_index sections appropriately
+        output['kv_cache_index'] = self._truncate_kv_cache_sections(
+            output['kv_cache_index'], self.max_length)
+        
+        return output
+
+    def _truncate_kv_cache_sections(self, kv_cache_sections: List[torch.Tensor], 
+                                  max_length: int) -> List[torch.Tensor]:
+        """Truncate kv_cache sections to fit within max_length."""
+        truncated_sections = []
+        current_pos = 0
+        
+        for section in kv_cache_sections:
+            section_length = section.size(1)
+            remaining_length = max_length - current_pos
+            
+            if remaining_length <= 0:
+                break
+            elif remaining_length >= section_length:
+                truncated_sections.append(section)
+                current_pos += section_length
+            else:
+                truncated_section = section[:, :remaining_length]
+                truncated_sections.append(truncated_section)
+                break
+        
+        return truncated_sections
+
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
+        """
+        Main collation function with improved logic.
+        
+        Args:
+            features: List of feature dictionaries from dataset
+            
+        Returns:
+            Batched and padded output dictionary
+        """
+        if not features:
+            return {}
+        
+        # Step 1: Normalize input format for all features
+        normalized_features = [self._normalize_input_format(feat) for feat in features]
+        
+        # Step 2: Split each feature into sections
+        all_sections = [self._split_into_sections(feat) for feat in normalized_features]
+        
+        # Step 3: Pad sections to create uniform batch structure
+        output = self._pad_sections(all_sections)
+        
+        # Step 4: Apply length constraints if needed
+        output = self._apply_length_constraints(output)
+        
+        return output
+
+
+class BaselineDataCollator:
+    """Custom data collator for baseline model training"""
+    
+    def __init__(self, tokenizer: AutoTokenizer, pad_to_multiple_of: Optional[int] = None):
+        self.tokenizer = tokenizer
+        self.pad_to_multiple_of = pad_to_multiple_of
+    
+    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+        # Extract input_ids and labels
+        input_ids = [f["input_ids"] for f in features]
+        labels = [f["labels"] for f in features]
+        
+        # Find max length in batch
+        max_length = max(len(ids) for ids in input_ids)
+        
+        # Apply pad_to_multiple_of if specified
+        if self.pad_to_multiple_of is not None:
+            max_length = ((max_length + self.pad_to_multiple_of - 1) // self.pad_to_multiple_of) * self.pad_to_multiple_of
+        
+        # Pad sequences
+        batch_input_ids = []
+        batch_labels = []
+        batch_attention_mask = []
+        
+        for ids, lbls in zip(input_ids, labels):
+            # Pad input_ids
+            padded_ids = ids + [self.tokenizer.pad_token_id] * (max_length - len(ids))
+            batch_input_ids.append(padded_ids)
+            
+            # Pad labels (use -100 for padding)
+            padded_labels = lbls + [-100] * (max_length - len(lbls))
+            batch_labels.append(padded_labels)
+            
+            # Create attention mask
+            attention_mask = [1] * len(ids) + [0] * (max_length - len(ids))
+            batch_attention_mask.append(attention_mask)
+        
+        return {
+            "input_ids": torch.tensor(batch_input_ids, dtype=torch.long),
+            "labels": torch.tensor(batch_labels, dtype=torch.long),
+            "attention_mask": torch.tensor(batch_attention_mask, dtype=torch.long),
+        }
+
+
+
+"""
+Helper functions
+"""
+
+
+def create_dataset(dataset_type: str, **kwargs) -> Dataset:
+    """
+    Factory function to create a dataset based on type.
+    
+    Args:
+        dataset_type: String indicating the type of dataset
+        **kwargs: Additional arguments to pass to the dataset constructor
+        
+    Returns:
+        An instance of the appropriate dataset
+    """
+    # First, check if dataset_type is directly in the registry (exact match)
+    if dataset_type in DATASET_REGISTRY:
+        return DATASET_REGISTRY[dataset_type](**kwargs)
+    
+    # Then check for case-insensitive match
+    dataset_type_lower = dataset_type.lower()
+    if dataset_type_lower in DATASET_REGISTRY:
+        return DATASET_REGISTRY[dataset_type_lower](**kwargs)
+    
+    # If not found in registry, raise an error with valid options
+    valid_options = list(
+        set([name for name, cls in DATASET_REGISTRY.items() if name == cls.__name__])
+    )  # Only include actual class names
+    raise ValueError(
+        f"Unknown dataset type: {dataset_type}. Valid options are: {valid_options}"
+    )

rosetta/train/model_utils.py

@@ -0,0 +1,142 @@
+"""
+Model setup utilities for RosettaModel training/evaluation
+"""
+
+import torch
+from typing import Dict, Any, List
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from rosetta.model.wrapper import RosettaModel
+from rosetta.model.projector import create_projector
+
+"""
+Mapping strategies
+"""
+def k_nearest_sources(num_target_layers: int, num_source_layers: int, k: int) -> Dict[int, List[int]]:
+    """
+    Compute a per-target mapping to K nearest source layers.
+
+    Returns: Dict[target_idx, List[source_idx]] only for targets we map.
+    Distances are computed by placing target and source layers uniformly in [0, 1]
+    and sorting by absolute distance.
+    """
+    if num_target_layers <= 1:
+        target_positions = [0.0]
+    else:
+        target_positions = [i / (num_target_layers - 1) for i in range(num_target_layers)]
+    if num_source_layers <= 1:
+        source_positions = [0.0]
+    else:
+        source_positions = [j / (num_source_layers - 1) for j in range(num_source_layers)]
+
+    mapping: Dict[int, List[int]] = {}
+    for t_idx, t_pos in enumerate(target_positions):
+        sorted_src = sorted(range(num_source_layers), key=lambda j: abs(source_positions[j] - t_pos))
+        chosen = sorted_src[:max(0, k)]
+        if len(chosen) > 0:
+            mapping[t_idx] = chosen
+    return mapping
+
+
+def last_aligned_sources(num_target_layers: int, num_source_layers: int, k: int = 1) -> Dict[int, List[int]]:
+    """
+    Return a per-target mapping that aligns the last target layer to the last
+    source layer and walks toward the front.
+
+    Returns: Dict[target_idx, List[source_idx]] only for targets we map. For each
+    target t, we choose up to K sources anchored at the aligned index, preferring
+    backward indices first then forward to satisfy K.
+
+    Example (T=11, S=33): target 10 -> [32, 31, ...], target 9 -> [31, 30, ...]
+    """
+    mapping: Dict[int, List[int]] = {}
+    if num_target_layers <= 0 or num_source_layers <= 0:
+        return mapping
+
+    # Align ends; offset >= 0 means extra source layers at the front
+    offset = num_source_layers - num_target_layers
+
+    def take_k_from(s0: int) -> List[int]:
+        result: List[int] = []
+        # Prefer moving backward from the anchor (last-to-front)
+        for back in range(k):
+            idx = s0 - back
+            if 0 <= idx < num_source_layers:
+                result.append(idx)
+        # If not enough due to boundary, extend forward
+        next_idx = s0 + 1
+        while len(result) < k and next_idx < num_source_layers:
+            result.append(next_idx)
+            next_idx += 1
+        return result
+
+    for t in range(num_target_layers):
+        s0 = offset + t
+        # Clamp to valid range for edge cases (e.g., fewer source layers)
+        if s0 < 0:
+            s0 = 0
+        elif s0 > num_source_layers - 1:
+            s0 = num_source_layers - 1
+        chosen = take_k_from(s0)
+        if len(chosen) > 0:
+            mapping[t] = chosen
+
+    return mapping
+
+
+def setup_models(model_config: Dict[str, Any], device: str = "cuda", dtype: torch.dtype = torch.bfloat16):
+    """Setup RosettaModel with base model, teacher model, and projectors"""
+    
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(model_config["base_model"])
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    # Load models
+    base_model = AutoModelForCausalLM.from_pretrained(
+        model_config["base_model"],
+        torch_dtype=dtype,
+        device_map=device
+    )
+    
+    teacher_model = AutoModelForCausalLM.from_pretrained(
+        model_config["teacher_model"],
+        torch_dtype=dtype,
+        device_map=device
+    )
+    
+    # Create projector
+    projector_config = model_config["projector"]
+    projector_params = projector_config["params"].copy()
+    projector_params["dtype"] = dtype
+    
+    projector = create_projector(
+        projector_config["type"],
+        source_dim=teacher_model.config.head_dim,
+        target_dim=base_model.config.head_dim,
+        **projector_params
+    )
+
+    # Setup RosettaModel
+    rosetta_model = RosettaModel(
+        model_list=[base_model, teacher_model],
+        base_model_idx=0,
+        projector_list=[projector]
+    ).to(device)
+    
+    # Configure projector mappings
+    num_layers_to_map = min(
+        base_model.config.num_hidden_layers, 
+        teacher_model.config.num_hidden_layers
+    )
+    
+    for layer_idx in range(num_layers_to_map):
+        rosetta_model.set_projector_config(
+            source_model_idx=1,  # Teacher
+            source_model_layer_idx=layer_idx,
+            target_model_idx=0,  # Base
+            target_model_layer_idx=layer_idx,
+            projector_idx=0
+        )
+    
+    return rosetta_model, tokenizer 

rosetta/utils/evaluate.py

@@ -0,0 +1,639 @@
+"""
+Common evaluation utilities for benchmark tasks.
+
+This module provides shared functions for model evaluation across different benchmarks
+like MMLU-Redux and MMMLU.
+"""
+
+import re
+import os
+import torch
+import torch.nn as nn
+import numpy as np
+from typing import Dict, Any, List, Tuple, Optional
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from rosetta.model.projector import load_projector
+from rosetta.model.wrapper import RosettaModel
+from rosetta.model.oracle import OracleRosettaModel
+
+def build_prompt(dataset: str, locale: str, question: str, choices: str, use_cot: bool, use_template: bool = True) -> str:
+    """
+    Build a localized prompt for a given dataset and locale.
+
+    Currently supports:
+    - dataset: "mmmlu"
+      - locale: "SW_KE" (Swahili). Other locales fall back to English.
+
+    Args:
+        dataset: Dataset identifier (e.g., "mmmlu")
+        locale: Locale/subject code (e.g., "SW_KE")
+        question: Question text
+        choices: Formatted choices string
+        use_cot: Whether to include CoT instruction
+
+    Returns:
+        Localized prompt string
+    """
+    
+        # Unified default English templates (shared by MMLU and MMMLU)
+    if not use_cot:
+        template = """Accurately answer the following question:
+
+{{question}}
+
+Choices:
+{{choices}}
+
+Instructions:
+- Carefully read the question and all options.
+- Select the single most correct answer.
+- Respond ONLY in the following format: "The correct answer is A/B/C/D".
+- Do not include any explanations, additional text, or punctuation besides the answer.
+
+The correct answer is"""
+
+    else:
+        template = """Accurately answer the following question:
+                   
+{{question}}
+
+Choices:
+{{choices}}
+
+Instructions:
+- Carefully read the question and all options.
+- Let's think step by step and explain your reasoning briefly.
+- Then give the final answer starting with The correct answer is"""
+
+    prompt = template.replace("{{question}}", question)
+    prompt = prompt.replace("{{choices}}", choices)
+
+    if not use_template:
+        prompt = question + "\n\nChoices:\n" + choices
+
+    return prompt
+
+
+def parse_answer(answer_str: str) -> List[str]:
+    """
+    Parse answer string to extract valid answer options.
+    Converts digits 0/1/2/3 to letters A/B/C/D.
+    
+    Args:
+        answer_str: String containing answer digits
+        
+    Returns:
+        List of parsed answer letters
+    """
+    if not isinstance(answer_str, str):
+        return []
+    valid_digits = [c for c in answer_str if c in {'0','1','2','3'}]
+    return sorted(list({
+        chr(65 + int(d))  # 0->A, 1->B, 2->C, 3->D
+        for d in valid_digits
+    }))
+
+
+def extract_answer_from_content(text: str) -> Optional[str]:
+    """
+    Extract answer from model output with robust multi-pattern matching.
+    Supports multiple languages and response formats.
+    
+    Args:
+        text: Model output text
+        
+    Returns:
+        Extracted answer letter or None
+    """
+    text = text.strip()
+    if not text:
+        return None
+
+    # Define multiple answer patterns for different languages and formats
+    answer_patterns = [
+        # English patterns
+        r'Answer:\s*(.*)',
+        r'answer:\s*(.*)',
+        r'ANSWER:\s*(.*)',
+        r'Your answer:\s*(.*)',
+        r'your answer:\s*(.*)',
+        r'YOUR ANSWER:\s*(.*)',
+        r'The answer is\s*(.*)',
+        r'the answer is\s*(.*)',
+        r'THE ANSWER IS\s*(.*)',
+        r'Correct answer is\s*(.*)',
+        r'correct answer is\s*(.*)',
+        r'Correct answer is:\s*(.*)',
+        r'correct answer is:\s*(.*)',
+        r'Correct answer:\s*(.*)',
+        r'correct answer:\s*(.*)',
+        r'CORRECT ANSWER:\s*(.*)',
+        
+        # Swahili patterns
+        r'Jibu lako:\s*(.*)',
+        r'jibu lako:\s*(.*)',
+        r'JIBU LAKO:\s*(.*)',
+        r'Jibu:\s*(.*)',
+        r'jibu:\s*(.*)',
+        r'JIBU:\s*(.*)',
+        r'Jibu sahihi:\s*(.*)',
+        r'jibu sahihi:\s*(.*)',
+        r'JIBU SAHIHI:\s*(.*)',
+        
+        # Other common patterns
+        r'Response:\s*(.*)',
+        r'response:\s*(.*)',
+        r'RESPONSE:\s*(.*)',
+        r'Choice:\s*(.*)',
+        r'choice:\s*(.*)',
+        r'CHOICE:\s*(.*)',
+        r'Option:\s*(.*)',
+        r'option:\s*(.*)',
+        r'OPTION:\s*(.*)',
+    ]
+    
+    # 1. Try to match any of the answer patterns
+    for pattern in answer_patterns:
+        match = re.search(pattern, text, re.IGNORECASE)
+        if match:
+            answer_part = match.group(1).strip()
+            # Search for first A-D letter in the matched part
+            for char in answer_part:
+                if char in {'A', 'B', 'C', 'D'}:
+                    return char
+    
+    # 2. Look for standalone A-D letters that are likely answers
+    # Prioritize letters at the end of text or with clear answer-like context
+    standalone_patterns = [
+        r'\b([A-D])(?:\s*[.,!?:)]?\s*$)',  # A-D at end of text with optional punctuation
+        r'\b([A-D])(?:\s*[.,!?:)]\s)',     # A-D followed by punctuation and space
+        r'(?:^|\s)([A-D])(?:\s*$)',        # A-D at start or with word boundary at end
+    ]
+    
+    for pattern in standalone_patterns:
+        matches = re.findall(pattern, text, re.IGNORECASE)
+        if matches:
+            # Check if this looks like mathematical expressions rather than answers
+            math_indicators = ['+', '-', '*', '/', '=', '^', 'x^', 'y^', 'z^', 'mod', 'sqrt', 'sin', 'cos', 'tan']
+            has_math = any(indicator in text for indicator in math_indicators)
+            has_answer_indicators = any(phrase in text.lower() for phrase in ['jibu', 'answer', 'choice', 'option', 'response', 'correct', 'sahihi'])
+            
+            # If it has math indicators but no answer indicators, it's likely mathematical notation
+            if has_math and not has_answer_indicators:
+                continue  # Skip this match, try next pattern
+            
+            return matches[-1].upper()
+    
+    # 3. Fallback: find all A-D letters but be more selective
+    all_letters = re.findall(r'\b([A-D])\b', text, re.IGNORECASE)
+    if all_letters:
+        # Check if this looks like mathematical expressions rather than answers
+        math_indicators = ['+', '-', '*', '/', '=', '^', 'x^', 'y^', 'z^', 'mod', 'sqrt', 'sin', 'cos', 'tan']
+        has_math = any(indicator in text for indicator in math_indicators)
+        has_answer_indicators = any(phrase in text.lower() for phrase in ['jibu', 'answer', 'choice', 'option', 'response', 'correct', 'sahihi'])
+        
+        # If it has math indicators but no answer indicators, it's likely mathematical notation
+        if has_math and not has_answer_indicators:
+            return None
+        
+        # Otherwise, return the last letter found
+        return all_letters[-1].upper()
+    
+    # 3. Search backwards for any A-D letter as fallback
+    for char in reversed(text):
+        if char in {'A', 'B', 'C', 'D'}:
+            return char
+
+    return None
+
+
+def apply_generation_config(model: Any, generation_config: Optional[Dict[str, Any]] = None) -> None:
+    """
+    Apply generation configuration to a model and handle sampling parameters.
+    
+    This function applies the provided generation config to the model and removes
+    sampling parameters (temperature, top_p, top_k, min_p) when do_sample=False
+    to avoid warnings from the transformers library. If no config is provided,
+    it defaults to greedy decoding with cleaned sampling parameters.
+    
+    Args:
+        model: Model object with generation_config attribute
+        generation_config: Optional generation configuration dictionary.
+                          If None, defaults to greedy decoding (do_sample=False).
+    """
+    if not hasattr(model, 'generation_config'):
+        return
+    
+    # If no config provided, default to greedy decoding
+    if not generation_config:
+        generation_config = {'do_sample': False}
+    
+    # Apply all configuration parameters
+    for key, value in generation_config.items():
+        setattr(model.generation_config, key, value)
+    
+    # Disable sampling parameters if do_sample=False to avoid warnings
+    # We set them to None instead of deleting, since some model code may
+    # access these attributes unconditionally.
+    if not generation_config.get('do_sample', True):
+        sampling_params = ['temperature', 'top_p', 'top_k', 'min_p', 'repetition_penalty']
+        for param in sampling_params:
+            try:
+                setattr(model.generation_config, param, None)
+            except Exception:
+                # If the backend does not allow setting, ignore silently
+                pass
+
+
+def set_default_chat_template(tokenizer, model_name: str):
+    """
+    Set default chat template for models without one.
+    
+    Args:
+        tokenizer: Tokenizer object
+        model_name: Name of the model
+    """
+    if tokenizer.chat_template is None:
+        if "UlizaLlama3".lower() in model_name.lower():
+            tokenizer.chat_template = (
+                "{%- for message in messages %}"
+                "{{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\\n\\n' }}"
+                "{{- message['content'] }}"
+                "{{- '<|eot_id|>' }}"
+                "{%- endfor %}"
+                "{%- if add_generation_prompt %}"
+                "{{- '<|start_header_id|>assistant<|end_header_id|>\\n\\n' }}"
+                "{%- endif %}"
+            )
+        else:
+            print(f"Model {model_name} has no chat template, setting default template...")
+            default_template = """{% for message in messages %}{% if message['role'] == 'user' %}### Human: {{ message['content'] }}{% elif message['role'] == 'assistant' %}### Assistant: {{ message['content'] }}{% endif %}{% if not loop.last %}
+    {% endif %}{% endfor %}{% if add_generation_prompt %}
+    ### Assistant:{% endif %}"""
+            tokenizer.chat_template = default_template
+            print("Default chat template has been set.")
+    else:
+        print(f"Model {model_name} already has a chat template.")
+
+
+def load_hf_model(model_name: str, device: torch.device, generation_config: Optional[Dict[str, Any]] = None) -> Tuple[Any, Any]:
+    """
+    Load Hugging Face model and tokenizer.
+    
+    Args:
+        model_name: Model name or path
+        device: Device to load model on
+        generation_config: Optional generation configuration to apply
+        
+    Returns:
+        Tuple of (model, tokenizer)
+    """
+    tokenizer = AutoTokenizer.from_pretrained(
+        str(model_name),
+        trust_remote_code=True,
+        padding_side='left'
+    )
+
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # Check and set chat template
+    set_default_chat_template(tokenizer, model_name)
+
+    if model_name == "google/gemma-3-1b-it":
+        torch._dynamo.config.cache_size_limit = 64
+        model = AutoModelForCausalLM.from_pretrained(
+            str(model_name),
+            torch_dtype=torch.bfloat16,
+            device_map={"": device},
+            sliding_window=4096
+        ).eval()
+    else:
+        model = AutoModelForCausalLM.from_pretrained(
+            str(model_name),
+            torch_dtype=torch.bfloat16,
+            device_map={"": device}
+    ).eval()
+    
+    # Apply generation config
+    apply_generation_config(model, generation_config)
+    
+    return model, tokenizer
+
+
+def load_rosetta_model(model_config: Dict[str, Any], eval_config: Dict[str, Any], 
+                      device: torch.device, generation_config: Optional[Dict[str, Any]] = None) -> Tuple[Any, Any]:
+    """
+    Load Rosetta model with projectors and aggregators.
+    
+    Args:
+        model_config: Model configuration dict
+        eval_config: Evaluation configuration dict
+        device: Device to load model on
+        generation_config: Optional generation configuration to apply
+        
+    Returns:
+        Tuple of (rosetta_model, tokenizer)
+    """
+    # Prefer checkpoints_dir under model.rosetta_config; fall back to eval config for backward compatibility
+    rosetta_config = model_config["rosetta_config"]
+    checkpoint_dir = rosetta_config.get("checkpoints_dir", eval_config.get("checkpoints_dir"))
+    if checkpoint_dir is None:
+        raise KeyError("checkpoints_dir must be provided under model.rosetta_config (preferred) or eval config (legacy)")
+    slm_model_path = rosetta_config["base_model"]
+    llm_model_path = rosetta_config["teacher_model"]
+
+    # Load tokenizer
+    slm_tokenizer = AutoTokenizer.from_pretrained(str(slm_model_path))
+    set_default_chat_template(slm_tokenizer, slm_model_path)
+    
+    # Load models
+    slm_model = AutoModelForCausalLM.from_pretrained(
+        str(slm_model_path),
+        torch_dtype=torch.bfloat16,
+        device_map={"": device}
+    ).eval()
+    
+    # Apply generation config to SLM
+    apply_generation_config(slm_model, generation_config)
+    
+    if llm_model_path == "google/gemma-3-1b-it":
+        llm_model = AutoModelForCausalLM.from_pretrained(
+            str(llm_model_path),
+                torch_dtype=torch.bfloat16,
+                device_map={"": device},
+                sliding_window=4096
+            ).eval()
+    else:
+        llm_model = AutoModelForCausalLM.from_pretrained(
+            str(llm_model_path),
+            torch_dtype=torch.bfloat16,
+            device_map={"": device}
+        ).eval()
+    
+    # Apply generation config to LLM
+    apply_generation_config(llm_model, generation_config)
+    
+    # Load projectors
+    num_projectors = len([f for f in os.listdir(checkpoint_dir) if re.match(r"projector_\d+\.pt", f)])
+    projector_list = []
+    for t in range(num_projectors):
+        json_cfg = os.path.join(checkpoint_dir, f"projector_{t}.json")
+        proj = load_projector(json_cfg)
+        proj = proj.to(device)
+        pt_path = os.path.join(checkpoint_dir, f"projector_{t}.pt")
+        if os.path.exists(pt_path):
+            state_dict = torch.load(pt_path, map_location=device)
+            proj.load_state_dict(state_dict, strict=False)
+        projector_list.append(proj)
+    
+    aggregator_list = []
+    
+    # Initialize Rosetta model
+    rosetta_model = RosettaModel(
+        model_list=[slm_model, llm_model],
+        base_model_idx=0,
+        projector_list=projector_list,
+        aggregator_list=aggregator_list,
+    ).to(device).eval()
+
+    # Load projector/aggregator mapping configs
+    proj_cfg_path = os.path.join(checkpoint_dir, "projector_config.json")
+    agg_cfg_path = os.path.join(checkpoint_dir, "aggregator_config.json")
+    rosetta_model.load_projector_config(proj_cfg_path)
+    rosetta_model.load_aggregator_config(agg_cfg_path)
+
+    return rosetta_model, slm_tokenizer
+
+
+def load_oracle_rosetta_model(model_config: Dict[str, Any], eval_config: Dict[str, Any], 
+                      device: torch.device) -> Tuple[Any, Any]:
+    """
+    Load Rosetta model with projectors and aggregators.
+    
+    Args:
+        model_config: Model configuration dict
+        eval_config: Evaluation configuration dict
+        device: Device to load model on
+        
+    Returns:
+        Tuple of (rosetta_model, tokenizer)
+    """
+    # Prefer checkpoints_dir under model.rosetta_config; fall back to eval config for backward compatibility
+    rosetta_config = model_config["rosetta_config"]
+    checkpoint_dir = rosetta_config.get("checkpoints_dir", eval_config.get("checkpoints_dir"))
+    if checkpoint_dir is None:
+        raise KeyError("checkpoints_dir must be provided under model.rosetta_config (preferred) or eval config (legacy)")
+    slm_model_path = rosetta_config["base_model"]
+    llm_model_path = rosetta_config["teacher_model"]
+
+    # Load tokenizer
+    slm_tokenizer = AutoTokenizer.from_pretrained(str(slm_model_path))
+    set_default_chat_template(slm_tokenizer, slm_model_path)
+    
+    # Load models
+    slm_model = AutoModelForCausalLM.from_pretrained(
+        str(slm_model_path),
+        torch_dtype=torch.bfloat16,
+        device_map={"": device}
+    ).eval()
+    
+    llm_model = AutoModelForCausalLM.from_pretrained(
+        str(llm_model_path),
+        torch_dtype=torch.bfloat16,
+        device_map={"": device}
+    ).eval()
+    
+    # Load projectors
+    num_projectors = len([f for f in os.listdir(checkpoint_dir) if re.match(r"projector_\d+\.pt", f)])
+    projector_list = []
+    for t in range(num_projectors):
+        json_cfg = os.path.join(checkpoint_dir, f"projector_{t}.json")
+        proj = load_projector(json_cfg)
+        proj = proj.to(device)
+        pt_path = os.path.join(checkpoint_dir, f"projector_{t}.pt")
+        if os.path.exists(pt_path):
+            state_dict = torch.load(pt_path, map_location=device)
+            proj.load_state_dict(state_dict, strict=False)
+        projector_list.append(proj)
+    
+    aggregator_list = []
+    
+    # Initialize Rosetta model
+    rosetta_model = OracleRosettaModel(
+        model_list=[slm_model, llm_model],
+        base_model_idx=0,
+        projector_list=projector_list,
+        aggregator_list=aggregator_list,
+    ).to(device).eval()
+
+    # Load projector/aggregator mapping configs
+    proj_cfg_path = os.path.join(checkpoint_dir, "projector_config.json")
+    agg_cfg_path = os.path.join(checkpoint_dir, "aggregator_config.json")
+    rosetta_model.load_projector_config(proj_cfg_path)
+    rosetta_model.load_aggregator_config(agg_cfg_path)
+
+    return rosetta_model, slm_tokenizer
+
+
+def get_option_token_ids(tokenizer, num_options: int = 4) -> List[int]:
+    """
+    Get token IDs for options A, B, C, D (or more up to J).
+    
+    Args:
+        tokenizer: Tokenizer object
+        num_options: Number of options to get (default 4 for A-D, max 10 for A-J)
+        
+    Returns:
+        List of token IDs for options
+    """
+    # Limit to maximum of 10 options (A-J)
+    num_options = min(num_options, 10)
+    option_ids = []
+    for i in range(num_options):
+        letter = chr(65 + i)  # A=65, B=66, etc.
+        ids = tokenizer.encode(" " + letter, add_special_tokens=False)
+        option_ids.append(ids[0] if ids else tokenizer.eos_token_id)
+    return option_ids
+
+"""
+Deprecated
+"""
+
+@torch.no_grad()
+def generate_answer_with_logits(model, tokenizer, prompt: str, option_ids: List[int], 
+                               device: torch.device, model_type: str = "hf") -> Tuple[str, np.ndarray]:
+    """
+    Generate answer using logits method.
+    
+    Args:
+        model: Model object
+        tokenizer: Tokenizer object
+        prompt: Input prompt
+        option_ids: Token IDs for options A, B, C, D
+        device: Device to run on
+        model_type: Type of model ("rosetta", "qwen", or "hf")
+        
+    Returns:
+        Tuple of (predicted_answer, probabilities)
+    """
+    messages = [{
+        "role": "user",
+        "content": prompt
+    }]
+    
+    # Try to apply chat template
+    try:
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False if model_type == "qwen" else None
+        )
+    except Exception as e:
+        print(f"Failed to apply chat template for {model_type} model: {e}")
+        text = f"### Human: {prompt}\n### Assistant:"
+    
+    text += "The correct answer is"
+    input_ids = tokenizer(text, return_tensors="pt").to(device)['input_ids']
+    attention_mask = torch.ones(input_ids.shape, dtype=torch.long).to(device)
+    position_ids = attention_mask.long().cumsum(-1) - 1
+    
+    if model_type == "rosetta":
+        instruction_index = torch.tensor([1, 0], dtype=torch.long).repeat(input_ids.shape[1]-1, 1).unsqueeze(0).to(device)
+        response_index = torch.tensor([[-1, 0]], dtype=torch.long).unsqueeze(0)
+        outputs = model.forward(
+            input_ids=input_ids, 
+            attention_mask=attention_mask, 
+            position_ids=position_ids, 
+            kv_cache_index=[instruction_index, response_index]
+        )
+    else:
+        outputs = model(input_ids)
+    
+    logits = outputs.logits[0, -1]
+    option_logits = torch.tensor([
+        logits[option_ids[0]].item(),
+        logits[option_ids[1]].item(),
+        logits[option_ids[2]].item(),
+        logits[option_ids[3]].item()
+    ])
+    
+    probs = torch.nn.functional.softmax(option_logits, dim=0).numpy()
+    pred = chr(65 + np.argmax(probs))
+    return pred, probs
+
+
+@torch.no_grad()
+def generate_answer_with_generate(model, tokenizer, prompt: str, device: torch.device,
+                                 model_type: str = "hf") -> Tuple[str, np.ndarray, int, int, str]:
+    """
+    Generate answer using text generation method.
+    
+    Args:
+        model: Model object
+        tokenizer: Tokenizer object
+        prompt: Input prompt
+        device: Device to run on
+        model_type: Type of model ("rosetta" or "hf")
+        
+    Returns:
+        Tuple of (predicted_answer, probabilities, input_length, generation_length, generated_text)
+    """
+    messages = [{
+        "role": "user",
+        "content": prompt
+    }]
+    
+    # Apply chat template
+    try:
+        text = tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=True,
+            enable_thinking=False
+        )
+    except Exception as e:
+        print(f"Failed to apply chat template: {e}")
+        text = f"### Human: {prompt}\n### Assistant:"
+
+    # Prepare model input
+    inputs = tokenizer(text, return_tensors="pt").to(device)
+    
+    # Generation parameters
+    sampling_params = {
+        'do_sample': True,
+        'temperature': 0.7,
+        'top_p': 0.8,
+        'top_k': 20,
+        'min_p': 0.0,
+        'repetition_penalty': 1.2,
+        'max_new_tokens': 1024
+    }
+    
+    # Generate text
+    outputs = model.generate(
+        input_ids=inputs.input_ids,
+        attention_mask=inputs.attention_mask,
+        **sampling_params
+    )
+    
+    # Parse output
+    if isinstance(model, RosettaModel):
+        generated_ids = outputs[0]
+    else:
+        generated_ids = outputs[0][inputs.input_ids.shape[1]:]
+    content = tokenizer.decode(generated_ids, skip_special_tokens=True).strip("\n")
+    
+    # Extract answer
+    pred = extract_answer_from_content(content)
+    
+    # Return uniform distribution for generate method
+    probs = np.array([0.25, 0.25, 0.25, 0.25])
+
+    input_length = inputs.input_ids.shape[1]
+    gen_length = generated_ids.shape[0]
+
+    return pred, probs, input_length, gen_length, content
+

rosetta/utils/matheval.py

@@ -0,0 +1,455 @@
+import re
+import json
+import argparse
+import jsonlines
+from collections import defaultdict
+from openai import OpenAI
+from typing import Dict, Any
+from math_verify import parse, verify, LatexExtractionConfig, ExprExtractionConfig, StringExtractionConfig
+from latex2sympy2_extended import NormalizationConfig
+from transformers import AutoTokenizer
+import requests
+import os
+
+class MathEvaluator:
+    
+    def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+        raise NotImplementedError
+
+    def extract_after_think(self, text: str, truncate_length: int = 1000, finish_generation: bool = True) -> str:
+        pattern = r"</think>(.*)"
+        match = re.search(pattern, text, re.DOTALL)
+        return match.group(1).strip() if (match and finish_generation) else text[-truncate_length:]
+    
+    def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extracted_answer: str = "", finish_generation: bool = True) -> str:
+        raise NotImplementedError
+
+    def get_llm_judge_prompt_not_finished(self, solution_str: str, ground_truth: str, extracted_answer: str = "", finish_generation: bool = True) -> str:
+        return f"""Please determine whether the final answer in the model-generated response was already correctly derived early in the reasoning process, and that the subsequent content consists mainly of unnecessary verification, overthinking, or repetitive reasoning. If correct is derived early, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+Reference answer: {ground_truth}
+Model-generated response: {solution_str}
+""".strip()
+
+    def llm_judge(self, solution_str: str, ground_truth: str, extracted_answer: str = "", finish_generation: bool = True) -> bool:
+        global OPENAI_CLIENT, MODEL_NAME
+        def get_inputs(scene_description):
+            body = [
+                {"role": "user", "content": scene_description},
+            ]
+            return body
+
+        def run_api(inputs):
+            completion = OPENAI_CLIENT.chat.completions.create(
+                model=MODEL_NAME,
+                messages=inputs
+            )
+            return completion.choices[0].message.content.strip()
+        if finish_generation:
+            scene_description = self.get_llm_judge_prompt(solution_str, ground_truth, extracted_answer, finish_generation)
+        else:
+            scene_description = self.get_llm_judge_prompt_not_finished(solution_str, ground_truth, extracted_answer, finish_generation)
+        inputs = get_inputs(scene_description)
+        response = run_api(inputs)
+
+        return "YES" in response
+
+
+class AIMEEvaluator(MathEvaluator):
+    def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+        # if not ground_truth.startswith("$"):
+        #     ground_truth = f"${ground_truth}$"
+        gold = parse(
+            ground_truth,
+            extraction_config=[ExprExtractionConfig()],
+        )
+        answer = parse(
+            solution_str,
+            extraction_config=[
+                LatexExtractionConfig(
+                    normalization_config=NormalizationConfig(
+                        nits=False,
+                        malformed_operators=False,
+                        basic_latex=True,
+                        boxed="all",
+                        units=True,
+                    ),
+                    boxed_match_priority=0,
+                    try_extract_without_anchor=False,
+                ),
+                ExprExtractionConfig(),
+            ],
+            extraction_mode="first_match",
+        )
+        if len(answer) == 0:
+            return False, "No extracted answer"
+        else:
+            return verify(gold, answer), str(answer)
+
+    def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extract_answer: str = "", finish_generation: bool = True) -> str:
+        solution_str = self.extract_after_think(solution_str, finish_generation=finish_generation)
+        return f"""Please determine whether the final answer provided in the model-generated response is equivalent to the reference answer from a math question. The final answer may either be enclosed in \\boxed{{}} or appear after "Answer:". If they are equivalent, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+Model-generated answer: {solution_str}
+Reference answer: {ground_truth}""".strip()
+
+
+class GSM8KEvaluator(MathEvaluator):
+    def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+        # if not ground_truth.startswith("$"):
+        #     ground_truth = f"${ground_truth}$"
+        gold = parse(
+            ground_truth,
+            extraction_config=[ExprExtractionConfig()],
+        )
+        answer = parse(
+            solution_str,
+            extraction_config=[
+                LatexExtractionConfig(
+                    normalization_config=NormalizationConfig(
+                        nits=False,
+                        malformed_operators=False,
+                        basic_latex=True,
+                        boxed="all",
+                        units=True,
+                    ),
+                    boxed_match_priority=0,
+                    try_extract_without_anchor=False,
+                ),
+                ExprExtractionConfig(),
+            ],
+            extraction_mode="first_match",
+        )
+        if len(answer) == 0:
+            return False, "No extracted answer"
+        else:
+            return verify(gold, answer), str(answer)
+
+    def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extract_answer: str = "", finish_generation: bool = True) -> str:
+        solution_str = self.extract_after_think(solution_str, finish_generation=finish_generation)
+        return f"""Please determine whether the final answer provided in the model-generated response with rule-based extracted answer is equivalent to the reference answer from a math question. The final answer may either be enclosed in the \\boxed{{}} or appear after the "Answer:". If they are equivalent, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+
+1. The reference answer does not include percentage signs, units or time formats (e.g., am, pm), but the Model-generated answer may include them.
+For example, 1 is equivalent to 1 %, 1 kg, 1 am, 1 pm, 1:00 am, 1:00 pm, etc.
+Model-generated answer: 1%
+Reference answer: 1
+Your output: YES
+
+Model-generated answer: 1 kg
+Reference answer: 1
+Your output: YES
+
+Model-generated answer: 1:00 pm
+Reference answer: 1
+Your output: YES
+
+2. The reference answer only includes one single number, but the Model-generated answer may include multiple numbers.
+For example, 10 is equivalent to \\boxed{{(4, 6)}}, etc.
+Model-generated answer: 5, 5
+Reference answer: 10
+Your output: YES
+
+Model-generated answer: 4, 6
+Reference answer: 10
+Your output: YES
+
+Model-generated answer: 86, 42
+Reference answer: 128
+Your output: YES
+
+Now let's try a real example.
+Model-generated answer: {solution_str}
+Reference answer: {ground_truth}
+""".strip()
+
+
+class MATH500Evaluator(MathEvaluator):
+    def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+        if not ground_truth.startswith("$"):
+            ground_truth = f"${ground_truth}$"
+        gold = parse(
+            ground_truth,
+            extraction_config=[LatexExtractionConfig()],
+        )
+        answer = parse(
+            solution_str,
+            extraction_config=[
+                LatexExtractionConfig(
+                    normalization_config=NormalizationConfig(
+                        nits=False,
+                        malformed_operators=False,
+                        basic_latex=True,
+                        boxed="all",
+                        units=True,
+                    ),
+                    boxed_match_priority=0,
+                    try_extract_without_anchor=False,
+                ),
+                ExprExtractionConfig(),
+            ],
+            extraction_mode="first_match",
+        )
+        if len(answer) == 0:
+            return False, "No extracted answer"
+        else:
+            return verify(gold, answer), str(answer)
+    def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extract_answer: str = "", finish_generation: bool = True) -> str:
+        solution_str = self.extract_after_think(solution_str, finish_generation=finish_generation)
+        return f"""Please determine whether the final answer provided in the model-generated response is equivalent to the reference answer from a math question. The final answer may either be enclosed in \\boxed{{}} or appear after "Answer:". If they are equivalent, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+Model-generated answer: {solution_str}
+Reference answer: {ground_truth}""".strip()
+    
+class AMCEvaluator(MathEvaluator):
+    def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+        if not ground_truth.startswith("$"):
+            ground_truth = f"${ground_truth}$"
+        gold = parse(
+            ground_truth,
+            extraction_config=[LatexExtractionConfig()],
+        )
+        answer = parse(
+            solution_str,
+            extraction_config=[
+                LatexExtractionConfig(
+                    normalization_config=NormalizationConfig(
+                        nits=False,
+                        malformed_operators=False,
+                        basic_latex=True,
+                        boxed="all",
+                        units=True,
+                    ),
+                    boxed_match_priority=0,
+                    try_extract_without_anchor=False,
+                ),
+                ExprExtractionConfig(),
+            ],
+            extraction_mode="first_match",
+        )
+        if len(answer) == 0:
+            return False, "No extracted answer"
+        else:
+            return verify(gold, answer), str(answer)
+    def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extract_answer: str = "", finish_generation: bool = True) -> str:
+        solution_str = self.extract_after_think(solution_str, finish_generation=finish_generation)
+        return f"""Please determine whether the final answer provided in the model-generated response is equivalent to the reference answer from a math question. The final answer may either be enclosed in \\boxed{{}} or appear after "Answer:". If they are equivalent, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+Model-generated answer: {solution_str}
+Reference answer: {ground_truth}""".strip()
+
+
+class GPQAEvaluator(MathEvaluator):
+    def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+        # if not ground_truth.startswith("$"):
+        #     ground_truth = f"${ground_truth}$"
+        gold = parse(
+            ground_truth,
+            extraction_config=[StringExtractionConfig()],
+        )
+        answer = parse(
+            solution_str,
+            extraction_config=[
+                StringExtractionConfig(),
+            ]
+        )
+        if len(answer) == 0:
+            return False, "No extracted answer"
+        else:
+            return verify(gold, answer), str(answer)
+        
+    def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extract_answer: str = "", finish_generation: bool = True) -> str:
+        solution_str = self.extract_after_think(solution_str, finish_generation=finish_generation)
+        return f"""Please determine whether the final answer provided in the model-generated response is equivalent to the reference answer from a multiple choice question. The final answer may either be enclosed in \\boxed{{}} or appear after "Answer:". If they are equivalent, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+Model-generated answer: {solution_str}
+Reference answer: {ground_truth}""".strip()
+
+
+# class MBPPEvaluator(Evaluator):
+#     def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+#         return True, "No extracted answer"
+        
+#     def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extract_answer: str = "", finish_generation: bool = True) -> str:
+#         solution_str = self.extract_after_think(solution_str, finish_generation=finish_generation)
+#         return f"""Please determine whether the final answer provided in the model-generated response is equivalent to the reference answer from a multiple choice question. The final answer may either be enclosed in \\boxed{{}} or appear after "Answer:". If they are equivalent, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+# Model-generated answer: {solution_str}
+# Reference answer: {ground_truth}""".strip()
+
+
+# class HUMANEVALEvaluator(Evaluator):
+#     def rule_judge(self, solution_str: str, ground_truth: str, finish_generation: bool = True) -> bool:
+#         return True, "No extracted answer"
+        
+#     def get_llm_judge_prompt(self, solution_str: str, ground_truth: str, extract_answer: str = "", finish_generation: bool = True) -> str:
+#         solution_str = self.extract_after_think(solution_str, finish_generation=finish_generation)
+#         return f"""Please determine whether the final answer provided in the model-generated response is equivalent to the reference answer from a multiple choice question. The final answer may either be enclosed in \\boxed{{}} or appear after "Answer:". If they are equivalent, return "YES"; if they are not, return "NO". Only return "YES" or "NO", and do not generate any other content.
+# Model-generated answer: {solution_str}
+# Reference answer: {ground_truth}""".strip()
+
+
+evaluator_map = {
+    "aime2024": AIMEEvaluator(),
+    "aime2025": AIMEEvaluator(),
+    "gsm8k": GSM8KEvaluator(),
+    "math500": MATH500Evaluator(),
+    "gpqa_diamond": GPQAEvaluator(),
+    "amc23": AMCEvaluator(),
+}
+
+API_BASE = None
+DEPLOYMENT_NAME = None
+API_VERSION = None
+CONSTRUCTED_URL = None
+API_KEY = None
+HEADERS = None
+OPENAI_CLIENT = None
+MODEL_NAME = None
+
+def set_client(api_base=None, deployment_name=None, api_version=None, api_key=None, model_name="gpt-4.1-2025-04-14"):
+    global API_BASE, DEPLOYMENT_NAME, API_VERSION, CONSTRUCTED_URL, API_KEY, HEADERS, MODEL_NAME, OPENAI_CLIENT
+
+    API_BASE = api_base
+    DEPLOYMENT_NAME = deployment_name
+    API_VERSION = api_version
+    CONSTRUCTED_URL = f"{api_base}/openai/deployments/{deployment_name}/chat/completions?api-version={api_version}"
+    API_KEY = api_key or os.getenv("OPENAI_API_KEY", "")
+    MODEL_NAME = model_name
+    HEADERS = {
+        "Content-Type": "application/json",
+        "api-key": api_key,
+    }
+    if API_KEY:
+        print(f"Using API key: {API_KEY}")
+        OPENAI_CLIENT = OpenAI(api_key=API_KEY)
+    else:
+        OPENAI_CLIENT = None
+    
+
+
+
+# def call_llm_judge(message: list, args: argparse.Namespace) -> str:
+#     """
+#     Call the Qwen API with the given message.
+    
+#     Args:
+#         message (list): Message list for the API.
+#         args (argparse.Namespace): Parsed arguments.
+        
+#     Returns:
+#         str: The content of the completion response.
+        
+#     Raises:
+#         Exception: When the API call fails.
+#     """
+
+#     try:
+#         completion = client.chat.completions.create(
+#             model=args.model_name,
+#             messages=message,
+#             temperature=args.temperature,
+#             top_p=args.top_p,
+#         )
+#         return completion.choices[0].message.content.strip()
+#     except Exception as e:
+#         print(f"API call failed: {str(e)}")
+#         raise
+
+
+# def rule_judge(completin, ground_truth, dataset_name: str) -> Dict[str, Any]:
+#     rule_judge_result = None
+
+#     rule_judge_result = evaluator_map[dataset_name].rule_judge(completin, ground_truth)
+
+#     if not rule_judge_result:
+#         print(f"No valid answer detected | LLM judge")
+#         call_llm_judge
+    
+        
+#     return example
+
+
+
+# def process_example(example: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]:
+#     """
+#     Process an individual example with token counting.
+#     """
+#     try:
+
+#         completion_text = example.get('completion', '')
+#         before_truncation_text = completion_text.split("</think>")[0]
+#         before_truncation_token_count = count_tokens(before_truncation_text, args)
+#         after_truncation_token_count = example.get('generated_tokens', 0) - before_truncation_token_count
+
+#         example['before_truncation_token_count'] = before_truncation_token_count
+#         example['after_truncation_token_count'] = after_truncation_token_count
+
+#         example = post_process(example)
+#         if not example.get('rule_judge_result', False):
+#             think_truncation = extract_after_think(completion_text) or ''
+
+#             prompt = generate_prompt(example, think_truncation)
+#             messages = format_prompt(prompt)
+#             llm_judge_response = call_qwen(messages, args)
+#             print(llm_judge_response)
+#             example['llm_judge_response'] = llm_judge_response
+#             if llm_judge_response == "YES":
+#                 example['llm_judge_result'] = True
+#                 example['final_judge_result'] = True
+#             elif llm_judge_response == "NO":
+#                 example['llm_judge_result'] = False
+#                 example['final_judge_result'] = False
+#             else:
+#                 example['llm_judge_result'] = None
+#                 example['final_judge_result'] = None
+                
+#         return example
+#     except Exception as e:
+#         print(f"Failed to process example {example.get('idx', 'unknown')}: {str(e)}")
+#         example['final_judge_result'] = None
+#         return example
+
+
+
+# def main():
+#     args = parse_arguments()
+#     try:
+#         final_results = []
+#         with jsonlines.open(args.result_save_name, mode='w') as writer:
+#             with open(args.data_load_name, 'r', encoding='utf-8') as f:
+#                 data = json.load(f)
+
+#             for idx, example in enumerate(data, start=1):
+#                 print(f"Processed example ID: {idx}")
+#                 new_example = process_example(example, args)
+#                 final_results.append(new_example)
+
+#             writer.write_all(final_results)
+#             print(f"Processing complete! Valid results saved to: {args.result_save_name}")
+
+#         grouped_data = defaultdict(list)
+#         for entry in final_results:
+#             grouped_data[entry['idx']].append(entry)
+
+#         pass_at_1_per_idx = {}
+#         for idx, entries in grouped_data.items():
+#             correct_count = sum(entry['final_judge_result'] for entry in entries if entry['final_judge_result'] is True)
+#             total_count = len(entries)
+#             pass_at_1_per_idx[idx] = correct_count / total_count if total_count else 0
+
+#         overall_pass_at_1 = sum(pass_at_1_per_idx.values()) / len(pass_at_1_per_idx)
+#         average_generated_tokens = sum(entry['generated_tokens'] for entry in final_results) / len(final_results)
+#         average_reasoning_tokens = sum(entry['before_truncation_token_count'] for entry in final_results) / len(final_results)
+#         average_answer_tokens = sum(entry['after_truncation_token_count'] for entry in final_results) / len(final_results)
+
+#         print(f'\nPass@1 per idx:\n{pass_at_1_per_idx}\n')
+#         print(f'Overall pass@1: {overall_pass_at_1:.4f}')
+#         print(f'Average generated_tokens: {average_generated_tokens:.2f}')
+#         print(f'Average reasoning_tokens: {average_reasoning_tokens:.2f}')
+#         print(f'Average answer_tokens: {average_answer_tokens:.2f}')
+
+#     except json.JSONDecodeError as e:
+#         print(f"Data loading failed: Invalid JSON format - {str(e)}")
+#         raise
+
+
+if __name__ == '__main__':
+    api_key = os.getenv("OPENAI_API_KEY", "")
+
+    set_client(api_base="", deployment_name="", api_version="", api_key=api_key)
+    response = "Okay, so I need to find the number of triples (a, b, c) of nonnegative integers such that a + b + c = 300 and the expression a\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b equals 6,000,000. Hmm, that seems a bit complicated, but maybe I can break it down step by step.\n\nFirst, let me note down the given equations:\n\n1. a + b + c = 300\n2. a\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b = 6,000,000\n\nI need to find all nonnegative integer triples (a, b, c) that satisfy both equations. Since the problem involves symmetric expressions in a, b, c, maybe there's a way to exploit symmetry here.\n\nLet me first if I can rewrite the second equation in a more manageable form. The expression a\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b looks symmetric. Maybe I can factor it somehow?\n\nLet me see: Each term is of the form (variable squared times another variable). So, for example, a\u00b2b + a\u00b2c = a\u00b2(b + c), similarly for the other terms. Let me try that:\n\na\u00b2(b + c) + b\u00b2(a + c) + c\u00b2(a + b)\n\nYes, that's the same as the original expression. Now, since a + b + c = 300, maybe I can substitute (b + c) with (300 - a), and similarly for the others. Let's try that substitution.\n\nSo, substituting:\n\na\u00b2(300 - a) + b\u00b2(300 - b) + c\u00b2(300 - c)\n\nWait, hold on. Let me check that again. The original expression is't exactly that. Wait, the's see:\n\nWait, the original expression is a\u00b2(b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b. So I group terms as a\u00b2(b + c) + b\u00b2(a + c) + c\u00b2(a + b), then yes, each term is variable squared times the sum of the other two variables. Since the sum of the other two variables is 300 minus the variable. So, for example, (b + c) = 300 - a, so substituting that in, we get:\n\na\u00b2(300 - a) + b\u00b2(300 - b) + c\u00b2(300 - c)\n\nTherefore, the second equation can be rewritten as:\n\na\u00b2(300 - a) + b\u00b2(300 - b) + c\u00b2(300 - c) = 6,000,000\n\nHmm, that might be a useful form. Let me write that down:\n\nEquation 2: a\u00b2(300 - a) + b\u00b2(300 - b) + c\u00b2(300 - c) = 6,000,000\n\nNow, perhaps I can consider expanding expression in terms of the variables. Let me think if there's a way to relate this to the first equation. Since a + b + c = 300, maybe we can express some relationship between these variables here Let't immediately yet.\n\nAlternatively, maybe I can consider expanding the expression or finding another way to express it. Let me see ifAlternatively, perhaps I can consider the expression a\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b. Let me see if this can be written in terms of (a + b + c)^3 or something similar.\n\nWait, let me recall that (a + b + c)^3 = a^3 + b^3 + c^3 + 3(a + b)(a + c)(b + c). Hmm, not sure if that helps. Alternatively, maybe expanding (a + b + c)(ab + bc + ca). Let me see:\n\n(a + b + c)(ab + bc + ca) = a^2b + a^2c + ab^2 + b^2c + a c^2 + b c^2 + 3abc\n\nWait, that's exactly our expression in the second equation plus 3abc. So:\n\n(a + b + c)(ab + bc + ca) = (a\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b) + 3abc\n\nTherefore, the second equation can be rewritten as:\n\n(a + b + c)(ab + bc + ca) - 3abc = 6,000,000\n\nSince a know that a + b + c = 300, substitute that in:\n\n300(abab + bc + ca) - 3abc = 6,000,000\n\nHmm, that might be a better way to handle it equation. Let me write that down:\n\nEqu00(ab + bc + ca) - 3abc = 6,000,000\n\nDiv can divide both sides by 3 to simplify:\n\n100(ab + bc + ca) - abc = 2,000,000\n\nSo now, the equation becomes:\n\n100(ab + bc + ca) - abc = 2,000,000\n\nHmm, that seems a bit more manageable. Let me write that as:\n\nabc - 100(ab + bc + ca) = -2,000,000\n\nAlternatively, maybe rearranged as:\n\nabc = 100(ab + bc + ca) - 2,000,000\n\nHmm, not sure yet. Let me see if I can find another relation here.\n\n Since we have a + b + c = 300, perhaps we can denote variables in terms of two variables. say, let me set variables as a, b, and then c = 300 - a - b. Then substitute expressions in terms of a and b. But that might get complicated, but perhaps manageable.\n\nAlternatively, maybe there's a substitution or symmetry here. Let me think.\n\n the problem is symmetric in a, b, c, so perhaps the solutions are symmetric in some way. Maybe two variables are equal, or all three are equal?\n\nLet me first check if all three variables are equal. If a = b = c, then since a + b + c = 300, each would be 100. Let me check if that satisfies the second equation.\n\nIf a = b = c = 100, then compute the second expression:\n\na\u00b2b + a\u00b2c + ...\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\n = 6*(100)^2*100 = 6*10003 = 6,000,000. Wait, that's exactly the required given! So (100, 100, 100) is a solution.\n\n So that's one solution. But the problem asks for all triples, so maybe might be more?\n\nHmm, so maybe there are other triples where variables are not all equal but still satisfy the equations. Let me see.\n\nAlternatively, maybe two variables are equal and the third is different. Let me suppose that two variables are equal, say a = b, and c is different. Let,2a + c = 300, so c = 300 - 2a.\n\nThen, let's plug into the second equation.\n\nFirst, let me compute ab + bc + ca. when a = b:\n\nab + bc + ca = a\u00b2 + a c + a c = a\u00b2 + 2a c\n\nSimilarly, abc = a\u00b2 c\n\nSo substitugging into the equation 100(ab + bc + ca) - abc = 100(a\u00b2 + 2a c) - a\u00b2 c = 100a\u00b2 + 200a c - a\u00b2 c\n\nBut since c = 300 - 2a, substitute that in:\n\n100a\u00b2 + 200a(300 - 2a) - a\u00b2(300 - 2a)\n\nLet me compute each term:\n\nFirst term: 100a\u00b2\n\nSecond term: 200a*(300 - 2a) = 200a*300 - 200a*2a = 60,000a - 400a\u00b2\n\nThird term: -a\u00b2*(300 - 2a) = -300a\u00b2 + 2a\u00b3\n\nSo adding all three terms together:\n\n100a\u00b2 + (60,000a - 400a\u00b2) + (-300a\u00b2 + 2a\u00b3)\n\nCombine like terms:\n\nFor00a\u00b2 - 400a\u00b2 - 300a\u00b2 + 2a\u00b3 + 60,000a\n\nThat's:\n\n(100 - 400 - 300)a\u00b2 + 2a\u00b3 + 60,000a\n\n= (-600a\u00b2) + 2a\u00b3 + 60,000a\n\nSo the entire expression is 2a\u00b3 - 600a\u00b2 + 60,000a\n\nSet this equal to 2,000,000 ( (since the equation was 100(ab + bc + ca) - abc = 2,000,000)\n\nWait:\n\n2a\u00b3 - 600a\u00b2 + 60,000a = 2,000,000\n\nDivide both sides by 2:\n\na\u00b3 - 300a\u00b2 + 30,000a = 1,000,000\n\nBring all terms to left-hand:\n\na\u00b3 - 300a\u00b2 + 30,000a - 1,000,000 = 0\n\nHmm, need to solve this cubic equation for integer a. Let me see if I can factor this or find integer roots.\n\nPossible me denote f(a) = a\u00b3 - 300a\u00b2 + 30,000a - 1,000,000\n\nWe can try possible integer roots. Since a must be a nonnegative integer such that 2a \u2264 300 (since c = 300 - 2a must be nonnegative). so a \u2264 150.\n\nPossible rational roots are factors of 1,000,000 divided by factors of 1, so possible integer roots are divisors of 1,000,000. But that's a lot. Maybe try some values.\n\nFirst, let me try a = 100. Since when all variables are 100, that's a solution. so if a = 100, then c = 300 - 200 = 100, so that's the case where all three are equal, so that should be a root.\n\nLet me check f(100):\n\n100\u00b3 - 300*(100)^2 + 30,000*100 - 1,000,000\n\n= 1,000,000 - 300*10,000 + 3,000,000 - 1,000,000\n\n= 1,000,000 - 3,000,000 + 3,000,000 - 1,000,000 = 0. Yes, so a = 100 is a root.\n\nTherefore, (a - 100) is a factor. Let's perform polynomial division or factor it out.\n\nDivide f(a) by (a - 100):\n\nUsing synthetic division:\n\nCoefficients: 1 | -300 | 30,000 | -1,000,000\n\nDivide by (a - 100), so root at 100.\n\nBring down the 1.\n\nMultiply 1 by 100: 100. Add to next coefficient: -300 + 100 = -200\n\nMultiply -200 by 100 = -20,000. Add to next coefficient: 30,000 + (-20,000) = 10,000\n\nMultiply 10,000 by 100 = 1,000,000. Add to last term: -1,000,000 + 1,000,000 = 0. Perfect.\n\nSo the polynomial factors as (a - 100)(a\u00b2 - 200a + 10,000)\n\nNow set the quadratic factor to zero:\n\na\u00b2 - 200a + 10,000 = 0\n\nLet me compute discriminant D = (200)^2 - 4*1*10,000 = 40,000 - 40,000 = 0\n\nSo discriminant is zero, so a = [200 \u00b1 0]/2 = 100. So the quadratic has a double root at a = 100.\n\nTherefore, the cubic factors as (a - 100)^3 = 0. Wait, but that would mean that the original cubic is (a - 100)^3, but let me check:\n\nWait, (a - 100)^3 = a\u00b3 - 300a\u00b2 + 30,000a - 1,000,000, which matches exactly the cubic. So indeed, the equation is (a - 100)^3 = 0, so the only real root is a = 100. Therefore, in the case where a = b, the only solution is a = b = 100, which gives c = 100. so that's the triple-equal case.\n\nHmm, so if we assume two variables are equal, the only solution is when all three are equal. So maybe the only symmetric with two variables equal is the all-equal case. So perhaps other are solutions where all variables are different?\n\nAlternatively, maybe there are solutions where one variable is 0. Let me check that possibility.\n\nSuppose c = 0. Then the equations become a + b = 300 and and the second equation becomes a\u00b2b + a\u00b2*0 + b\u00b2a + b\u00b2*0 + 0 + 0 = a\u00b2b + ab\u00b2 = ab(a + b) = ab*300.\n\nSo the second equation becomes 300ab = 6,000,000, so ab = 20,000.\n\nSo if have a + b = 300 and ab = 20,000. Let me see if there are integer solutions here here.\n\nWe can set up the quadratic equation: Let a and b be roots of x\u00b2 - 300x + 20,000 = 0.\n\nCompute discriminant D = 300\u00b2 - 4*1*20,000 = 90,000 - 80,000 = 10,000\n\nsqrt(D) = 100\n\n so roots are [300 \u00b1 100]/2 = (400)/2 = 200 or (200)/2 = 100. So the roots are 200 and 100. Therefore, the solutions are (a, b) = (200, 100) or (100, 200). Since c = 0, the triples would be (200, 100, 0), (100, 200, 0), and and also permutations where c is in different positions. Wait, since c was set to 0, but the variables are ordered, so actually, if we set c = 0, then the triples are (a, b, 0) where a + b = 300 and ab = 20,000. So the solutions are (200, 100, 0), (100, 200, 0). But since the variables are ordered, but the problem counts triples (a, b, c) as ordered? Wait, actually, the problem says \"triples of nonnegative integers (a, b, c)\", so I think order matters here. Wait, but actually, in combinatorics, unless specified otherwise, sometimes triples are considered ordered. But let me check the problem statement again.\n\nWait, the problem says \"the number of triples of nonnegative integers (a, b, c)\", so I think they are ordered triples. so different permutationsings count as different triples. However, in the case where c is 0, but we could also have a = 0 or b = 0. So actually's see.\n\nWait, so if c = 0, then we have two solutions (200,100,0 and 100,200,0. Similarly, if a = 0, then we have (0, 200, 100) and (0, 100, 200). Similarly, if b = 0, then (200, 0, 100) and (100, 0, 200). So total of 6 triples where one variable is 0 and the other two are 200 and 100. Additionally, the all-100 case. So that's 7 so far. But maybe there are more solutions where none of the variables are zero?\n\nAlternatively, maybe there are solutions where two variables are zero? Let me check that.\n\n quickly. Suppose two variables are zero, say b = c = 0, then a = 300. Then the second equation becomes a\u00b2b + ... all terms would be zero, so 0 \u2260 6,000,000. So that's invalid. So no solutions with two variables zero.\n\nSo the only solutions with a zero are the ones above with one zero and the other two being 200 and 100. So that gives 6 triples. Plus the all-100 case. So total 7? But maybe there are other triples where none of the variables are zero?\n\nWait, let me see.\n\n Let me suppose of another approach. Let me suppose variables a, b, c all positive integers (since if one is zero, we already considered that case). Let me see if there are solutions where all variables are positive and and different from 100, 100, 100 and 200,100,0 etc.\n\nAlternatively, maybe there can consider the equation 100(ab + bc + ca) - abc = 2,000,000. Let me denote this as:\n\nabc - 100(ab + bc + ca) = -2,000,000\n\nHmm, perhaps we can factor this equation somehow. Let me see:\n\nLet me add 1,000,000 to both sides:\n\nabc - 100(ab + bc + ca) + 1,000,000 = -1,000,000\n\nHmm, not sure. Alternatively, perhaps we't think of variables x = a, y = b, z = c, and see if there's a substitution.\n\nAlternatively, let set variables = 300 - a, y = 300 - b, z = 300 - c, but not sure.\n\nAlternatively, perhaps consider the equation:\n\n100(ab + bc + ca) - abc = 2,000,000. Let me rearr this as:\n\nabc - 100(ab - 100bc - 100ca + 2,000,000 = 0\n\nHmm, perhaps factor terms:\n\nLet me see, perhaps factor terms with ab, bc, ca:\n\nab(c - 100) + bc(a - 100) + ca(b - 100) = -2,000,000\n\nHmm, not sure. Alternatively, perhaps rearr the equation as:\n\nabc = 100(ab + bc + ca) - 2,000,000\n\nLet me think if I can write this as:\n\nabc = 100(ab + bc + ca) - 2,000,000\n\nHmm, perhaps divide both sides by abc (assuming a,b,c \u2260 zero):\n\n1 = 100(1/c + 1/a + 1/b) - 2,000,000/(abc)\n\nHmm, not sure if helpful helps.\n\nAlternatively, maybe think of variables in terms of ratios. Let me suppose that variables are in some proportion. Let me see.\n\nAlternatively, let me consider that in the case where one variable is 0, we have solutions. The case where all variables equal 100 is a solution. Maybe there are other solutions where variables are different but not involving zeros?\n\nAlternatively me suppose that variables a, b, c are all different from 100 and none are zero. Let me see if such solutions exist.\n\nAlternatively, let try to set variables as 150, 100, 50. Let me test that.\n\nWait, let me pick some numbers. Let me see, for't suppose a = 150, b = 100, c = 50. Then a + b + c = 300. Let me compute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\nCompute each term:\n\na\u00b2b = 150\u00b2 * 100 = 22500 * 100 = 2,250,000\n\na\u00b2c = 150\u00b2 *50 = 22500 *50 = 1,125,000\n\nb\u00b2a = 100\u00b2 *150 = 10,000 *150 = 1,500,000\n\nb\u00b2c = 100\u00b2 *50 = 10,000 *50 = 500,000\n\nc\u00b2a = 50\u00b2 *150 = 2500 *150 = 375,000\n\nc\u00b2b = 50\u00b2 *100 = 2500 *100 = 250,000\n\nAdding all these up:\n\n2,250,000 + 1,125,000 = 3,375,000\n\n+1,500,000 = 4,875,000\n\n+500,000 = 5,375,000\n\n+375,000 = 5,750,000\n\n+250,000 = 6,000,000\n\nWait a second! That works! So (150, 100, 50) is another solution. Hmm that's a solution where all variables are different and none are zero. So that's another solution. Hmm my previous thought that only the all-100 and the ones with a zero were solutions was incorrect. So there are more solutions.\n\nHmm, so that complicates things. Let me see why that works. Let me check the calculation again.\n\nWait, so (150, 100, 50) gives the required sum and the required product expression6,000,000. So that's a valid solution. So there are more solutions. So I need to find all such triples.\n\nHmm, so perhaps there are multiple families variables are in arithmetic progression? Let me see. 50, 100, 150 is an arithmetic progression with difference 50. Let me see if that's a pattern.\n\nAlternatively, maybe variables a can think of variables as multiples of 50? Let me see.\n\nWait, let me see if there are other triples. that. Let me try (200, 100, 0) which we already know works, but (150, 100, 50) works, and (100, 100, 100) works. Let me see if there are others.\n\nLet me try (200, 50, 50). Let's check:\n\na + b + c = 200 +50 +50 = 300. Good.\n\nCompute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\n= 200\u00b2*50 + 200\u00b2*50 + 50\u00b2*200 + 50\u00b2*50 + 50\u00b2*200 + 50\u00b2*50\n\nWait, let me compute each term:\n\na\u00b2b = (200)^2 *50 = 40,000 *50 = 2,000,000\n\na\u00b2c = same as a\u00b2b since b and c are both 50, so another 2,000,000\n\nb\u00b2a = (50)^2 *200 = 2,500 *200 = 500,000\n\nb\u00b2c = (50)^2 *50 = 2,500 *50 = 125,000\n\nc\u00b2a = same as b\u00b2a, so another 500,000\n\nc\u00b2b = same as b\u00b2c, so another 125,000\n\nAdding all together:\n\n2,000,000 + 2,000,000 = 4,000,000\n\n+500,000 + 00,000 = 4,625,000\n\nWait, no, let me do step by step:\n\nFirst two terms: 2,000,000 + 2,000,000 = 4,000,000\n\nThen b\u00b2a + b\u00b2c = 500,000 + 125,000 = 625,000\n\nThen c\u00b2a + c\u00b2b = same as b\u00b2a + b\u00b2c, so another 625,000\n\nTotal: 4,000,000 + 625,000 + 625,000 = 5,250,000. which is less than 6,000,000. So that doesn't work.\n\nHmm, so that's not a solution. So maybe arithmetic progression is a better idea.\n\nWait, let me try another triple. Let me see, maybe't try (120, 100, 80). Let me check sum: 120 + 100 +80 = 300. Good.\n\nCompute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\nCompute me compute each term:\n\na\u00b2b = 120\u00b2 *100 = 14,400 *100 = 1,440,000\n\na\u00b2c = 120\u00b2 *80 = 14,400 *80 = 1,152,000\n\nb\u00b2a = 100\u00b2 *120 = 10,000 *120 = 1,200,000\n\nb\u00b2c = 100\u00b2 *80 = 10,000 *80 = 800,000\n\nc\u00b2a = 80\u00b2 *120 = 6,400 *120 = 768,000\n\nc\u00b2b = 80\u00b2 *100 = 6,400 *100 = 640,000\n\nAdding all terms:\n\n1,440,000 + 1,152,000 = 2,592,000\n\n+1,200,000 = 3,792,000\n\n+800,000 = 4,592,000\n\n+768,000 = 5,360,000\n\n+640,000 = 6,000,000\n\nOh! That works too! So (120, 100, 80) is another solution. Interesting. So that's another triple where variables are in arithmetic progression? 80, 100, 120 with difference 20. Wait, so that's also an arithmetic progression. Hmm, so maybe arithmetic progressions are the key here?\n\nWait, let me see another arithmetic progression. Let me try (180, 100, 20). Let me check sumSum is 180 + 100 +20 = 300. Let me compute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\na\u00b2b = 180\u00b2 *100 = 32,400 *100 = 3,240,000\n\na\u00b2c = 180\u00b2 *20 = 32,400 *20 = 648,000\n\nb\u00b2a = 100\u00b2 *180 = 10,000 *180 = 1,800,000\n\nb\u00b2c = 100\u00b2 *20 = 10,000 *20 = 200,000\n\nc\u00b2a = 20\u00b2 *180 = 400 *180 = 72,000\n\nc\u00b2b = 20\u00b2 *100 = 400 *100 = 40,000\n\nAdding all terms:\n\n3,240,000 + 648,000 = 3,888,000\n\n+1,800,000 = 5,688,000\n\n+200,000 = 5,888,000\n\n+72,000 = 5,960,000\n\n+40,000 = 6,000,000\n\nWow, that works works! So (180, 100, 20) is another solution. So arithmetic progression with difference 80? Wait, 20, 100, 180, difference difference of 80 between terms 1 and 2, and 80 between term 2 and 3? Wait,100 -20 =80, 180 -100=80. Yes, so that's an arithmetic progression with common difference 80. So that works too as well.\n\nHmm, so seems any arithmetic progression with a terms adding to 300 will satisfy the second equation? Let me check another one.\n\n Let me try (160, 100, 40). Sum is 160 + 00 +40 = 300.\n\nCompute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\na\u00b2b = 160\u00b2 *100 = 25,600 *100 = 2,560,000\n\na\u00b2c = 160\u00b2 *40 = 25,600 *40 = 1,024,000\n\nb\u00b2a = 100\u00b2 *160 = 10,000 *160 = 1,600,000\n\nb\u00b2c = 100\u00b2 *40 = 10,000 *40 = 400,000\n\nc\u00b2a = 40\u00b2 *160 = 1,600 *160 = 256,000\n\nc\u00b2b = 40\u00b2 *100 = 1,600 *100 = 160,000\n\nAdding all terms:\n\n2,560,000 + ,024,000 = 3,584,000\n\n+1,600,000 = 5,184,000\n\n+400,000 = 5,584,000\n\n+256,000 = 5,840,000\n\n+160,000 = 6,000,000\n\nYes! That works too. So (160, 100, 40) is another solution. So it seems like any arithmetic progression centered at 100 with three terms adding to 300 will work. Let me see.\n\nWait, let me think of the arithmetic progression. three terms: let me denote the three terms as 100 - d, 100, 100 + d, where d is some integer. Then their sum is (100 - d) + 100 + (100 + d) = 300, which works. So that's an arithmetic progression with common difference d. Let me check if this satisfies the second equation.\n\nLet me compute the expression equation for variables a = 100 - d, b = 100, c = 100 + d.\n\nCompute a\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b.\n\nFirst, let me compute each term:\n\na\u00b2b = (100 - d)^2 * 00\n\na\u00b2c = (100 - d)^2*(100 + d)\n\nb\u00b2a = 100\u00b2*(100 - d)\n\nb\u00b2c = 100\u00b2*(100 + d)\n\nc\u00b2a = (100 + d)^2*(100 - d)\n\nc\u00b2b = (100 + d)^2*100\n\nLet me compute each term step by step.\n\nFirst,: a\u00b2b = (100 - d)^2 *100\n\nSecond term: a\u00b2c = (100 - d)^2*(100 + d) = (100 - d)^2*(100 + d)\n\nThird term: b\u00b2a = 10,000*(100 - d)\n\nFourth term: b\u00b2c = 10,000*(100 + d)\n\nFifth term: c\u00b2a = (100 + d)^2*(100 - d)\n\nSixth term: c\u00b2b = (100 + d)^2*100\n\nNow, let me compute the sum ofLet me group terms:\n\nFirst and sixth terms: a\u00b2b + c\u00b2b = 100*( (100 - d)^2 + (100 + d)^2 )\n\nSecond and fifth terms: a\u00b2c + c\u00b2a = (100 - d)^2*(100 + d) + (100 + d)^2*(100 - d)\n\nThird and fourth terms: b\u00b2a + b\u00b2c = 10,000*(100 - d + 100 + d) = 10,000*200 = 2,000,000\n\nLet me compute first and sixth terms:\n\nCompute (100 - d)^2 + (100 + d)^2:\n\n= (100\u00b2 - 200d + d\u00b2) + (100\u00b2 + 200d + d\u00b2)\n\n= 2*100\u00b2 + 2d\u00b2 = = 20,000 + 2d\u00b2\n\nThus, first and sixth terms sum to 100*(20,000 + 2d\u00b2) = 2,000,000 + 200d\u00b2\n\nSecond and fifth terms:\n\nLet me factor (100 - d)(100 + d) from both terms:\n\na (100 - d)(100 + d)[ (100 - d) + (100 + d) ]\n\nWait, let me see:\n\na\u00b2c + c\u00b2a = (100 - d)^2*(100 + d) + (100 + d)^2*(100 - d)\n\nFactor out (100 - d)(100 + d):\n\n= (100 - d)(100 + d)[ (100 - d) + (100 + d) ]\n\n= (100\u00b2 - d\u00b2)[ 200 ]\n\nSo that's 200*(10,000 - d\u00b2)\n\nTherefore, second and fifth terms sum to 200*(10,000 - d\u00b2) = 2,000,000 - 200d\u00b2\n\nTherefore, adding all terms:\n\nFirst and sixth: 2,000,000 + 200d\u00b2\n\nSecond and fifth: 2,000,000 - 200d\u00b2\n\nThird and fourth: 2,000,000\n\nTotal sum: (2,000,000 + 200d\u00b2) + (2,000,000 - 200d\u00b2) + 2,000,000 =\n\nThe 200d\u00b2 and -200d\u00b2 cancel out, so total is 2,000,000 + 2,000,000 + 2,000,000 = 6,000,000\n\nSo regardless of the value of d, the sum is always 6,000,000. That's fascinating! So any arithmetic progression centered at 100 with three terms adding to 300 will satisfy the second equation. Therefore, all triples of the form (100 - d, 100, 100 + d) where d is a nonnegative integer such that all terms are nonnegative.\n\nWait, but d can be positive or negative? Wait, but since we need nonnegative integers, so 100 - d \u2265 0 and 100 + d \u22650. Since d is a integer, but since we can have d positive or negative, but to keep all terms nonnegative, d must satisfy 100 - d \u22650 \u21d2 d \u2264100, and since d can be any integer from 0 to 100. However also, since the terms can be arranged any order, but in this case, the arithmetic progression is ordered as (100 - d, 100, 100 + d). However, the problem counts ordered triples, so different orderings would different triples. Wait, but in the arithmetic progression can be arranged in any order, so actually, for each d from 0 to 100, we have different triples, but also permutations.\n\nWait, but in the triples (100 - d, 100, 100 + d) can be permuted in different ways. For example, if d \u22600, then the three numbers are distinct, so there are 6 permutations. However, when d =0, all three are 100, so only one permutation.\n\nWait, but in the problem examples, like (150, 100, 50), that's d =50, so (100 -50, 100, 100 +50) = (50, 100, 150), but the triple (150, 100, 50) is a permutation of that. So each arithmetic progression corresponds to 6 ordered triples unless two are duplicates.\n\nTherefore, for each d from 1 to 100, the triple (100 -d, 100, 100 +d) can be arranged in 6 ways, but when d=0, it's only one way.\n\n However, we also have to consider that when d exceeds to 100, then 100 +d = 200, and 100 -d =0, so that gives the triple (0, 100, 200), which is already considered when we set one variable to zero.\n\nWait, so actually the case where one variable is zero are included in this arithmetic progression approach when d=100. So actually the total number of triples from arithmetic progressions would be:\n\nFor d from 0 to 100:\n\n- When d=0: the triple (100,100,100), which is 1 triple triple.\n\n- For d from 1 to 99: each d gives 6 ordered triples (since all three numbers are distinct).\n\n- When d =100: the triple (0, 100, 200), which can permutations in 3! /1! =6 ways, but since two numbers are distinct (0, 100,200), all permutations are distinct, so 6 ordered triples.\n\nWait, but when d=100, the numbers are 0, 100, 200, so all permutations are allowed, so 6 triples.\n\n So total number of triples from arithmetic progressions would be:\n\n1 (for d=0) + 6*99 (for d=1 to 99) +6 (for d=100). = 1 + 594 +6 = 601? Wait, but that can't be right because when d=100, it's included in the d=1 to 99? Wait no, d=100 is separate.\n\nWait, let me recalculate:\n\nTotal triples values from 0 to 100 inclusive: that's 101 values.\n\nFor each d from 1 to 99: 99 values, each gives 6 triples.\n\nFor d=0: 1 triple.\n\nFor d=100: 1 value, giving 6 triples.\n\nTotal triples: 1 + 99*6 +6 = 1 + 594 +6 = 601.\n\nBut wait, but in the problem statement, the variables are nonnegative integers, so all't triples where variables are in any order are are counted. separately separate. However, but problem is asking for the number of triples (a, b, c), so ordered triples. So this approach would give 601 triples. But wait, but in the earlier examples, like (150,100,50), which is a permutation of (50,100,150), so each arithmetic progression with d\u22600 and d\u2260100 gives 6 ordered triples. However, d=100, the triple is (0,100,200), which also has 6 permutations. And d=0 gives only one triple.\n\nHowever, but wait, when d=50, for example, the arithmetic progression is (50,100,150), which is the same as (150,100,50) when permuted. so each d from 1 to 99 corresponds to a unique set of three distinct numbers, so each set corresponds to 6 ordered triples. So the total number would be 1 + 6*100 = 601? Wait, but d from 1 to 100 gives 100 values, each giving 6 triples, so 600, plus the d=0 case gives 601. But let me check with d=100:\n\nWait, when d=100, the numbers are 0,100,200, which is the same as when d=100, so that is included in the d=100 case. So yes, that's correct.\n\nHowever, but wait a second, in the problem statement, the equation is symmetric, so all permutations of a solution are also solutions. Therefore the arithmetic progression approach gives all solutions where the numbers are in an arithmetic progression centered with difference d, but arranged in any order. So the total number of ordered triples would indeed 601. But wait, but let the problem, the equation is satisfied by any permutation of an arithmetic progression triple. so all permutations permutations are valid solutions.\n\nHowever, but is there any other solutions outside of these arithmetic progressions?\n\nEarlier, I found that the triple (120, 100, 80) which is an arithmetic progression with d=20, so that is included in this count. Similarly, (180,100,20) is d=80, so also included. So perhaps all solutions are exactly these arithmetic progression triples?\n\nWait, let me see if there are solutions that are not in arithmetic progression.\n\nSuppose I take a triple like (140, 100, 60). Let me check if that works.\n\nSum is 140 +100 +60 = 300.\n\nCompute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\na\u00b2140, b=100, c=60.\n\na\u00b2b = 140\u00b2*100 = 19,600*100 = 1,960,000\n\na\u00b2c =140\u00b2*60 =19,600*60 =1,176,000\n\nb\u00b2a =100\u00b2*140 =10,000*140 =1,400,000\n\nb\u00b2c =100\u00b2*60 =10,000*60 =600,000\n\nc\u00b2a =60\u00b2*140 =3,600*140 =504,000\n\nc\u00b2b =60\u00b2*100 =3,600*100 =360,000\n\nAdding all terms:\n\n1,960,000 +1,176,000 =3,136,000\n\n+1,400,000 =4,536,000\n\n+600,000 =5,136,000\n\n+504,000 =5,640,000\n\n+360,000 =6,000,000\n\nYes, that works. too. And ( this is an arithmetic progression with d=40: 60,100,140. So yes is included in the arithmetic progression case.\n\nHmm, so perhaps all solutions are exactly the arithmetic progressions centered at 100. Let me see if there is a solution that is not an arithmetic progression.\n\nLet me try a triple like (130, 110, 60). Sum is 130+110+60=300.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\na=130, b=110, c=60.\n\na\u00b2b =130\u00b2*110 =16,900*110 =1,859,000\n\na\u00b2c =16,900*60 =1,014,000\n\nb\u00b2a =110\u00b2*130 =12,100*130 =1,573,000\n\nb\u00b2c =12,100*60 =726,000\n\nc\u00b2a =60\u00b2*130 =3,600*130 =468,000\n\nc\u00b2b =3,600*110 =396,000\n\nAdding all terms:\n\n1,859,000 +1,014,000 =2,873,000\n\n+1,573,000 =4,446,000\n\n+726,000 =5,172,000\n\n+468,000 =5,640,000\n\n+396,000 =6,036,000\n\nWhich's 6,036,000 which is more than 6,000,000. So that doesn't work. So that triple is not a solution.\n\nHmm test: Let me try (120, 90, 150). Wait, sum is 160+90+150=400, which is over. Not good.\n\nLet me try (120, 90, 90). Sum is 120+90+90=300.\n\nCompute the second equation:\n\na\u00b22b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\na=120, b=90, c=90.\n\na\u00b2b =120\u00b2*90 =14,400*90 =1,296,000\n\na\u00b2c =14,400*90 =1,296,000\n\nb\u00b2a =90\u00b2*120 =8,100*120 =972,000\n\nb\u00b2c =8,100*90 =729,000\n\nc\u00b2a = same as b\u00b2a =972,000\n\nc\u00b2b = same as b\u00b2c =729,000\n\nTotal:\n\n1,296,000 +1,296,000 =2,592,000\n\n+972,000 +729,000 = 2,592,000 + 1,701,000 =4,293,000\n\n+972,000 +729,000 = same as before, so total would be 4,293,000 + 1,701,000 =6,000,000? Wait, let me compute step by step:\n\nWait, let me add all terms:\n\nFirst.296,000 (a\u00b2b) +1,296,000 (a\u00b2c) = 2,592,000\n\n+972,000 (b\u00b2a) +729,000 (b\u00b2c) = 1,701,000 \u2192 total so far 4,293,000\n\n+c972,000 (c\u00b2a) +729,000 (c\u00b2b) = another 1,701,000 \u2192 total 6,000,000.\n\nWait, so that works! So (120, 90, 90) is a solution. But this is not an arithmetic progression. Let me check:\n\nThe numbers are 120, 90, 90. So two of them are equal. So this is a case where two variables are equal, but not the middle term. Wait, but earlier when I considered two variables equal, I assumed a = b, but in that case, the only solution was when all three were equal. But here, with't two variables equal (b and c), but a different.\n\n from them. So this is a case where two variables are equal but not the first two.\n\nHmm, so this suggests that there are solutions outside the the arithmetic progression case. So my previous conclusion was wrong.\n\nWait, let let me analyze this case. Let me suppose that two variables are equal, say b = c, then a + 2b = 300, so a = 300 - 2b.\n\nThen, let's compute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\nSince b = c, this becomes:\n\na\u00b2b + a\u00b2b + b\u00b2a + b\u00b2b + b\u00b2a + b\u00b2b\n\n= 2a\u00b2b + 2b\u00b2a + 2b\u00b3\n\n= 2a\u00b2b + 2ab\u00b2 + 2b\u00b3\n\nFactor out 2b:\n\n2b(a\u00b2 + ab + b\u00b2)\n\nSet this equal to 6,000,000:\n\n2b(a\u00b2 + ab + b\u00b2) = 6,000,000\n\nDivide both sides by 2:\n\nb(a\u00b2 + ab + b\u00b2) = 3,000,000\n\nBut since a = 300 - 2b, substitute that in:\n\nb[(300 - 2b)^2 + (300 - 2b)b + b\u00b2] = 3,000,000\n\nLet me compute the expression inside the brackets:\n\nFirst term: (300 - 2b)^2 = 90,000 - 1,200b +4b\u00b2\n\nSecond term: (300 - 2b)bb = 300b - 2b\u00b2\n\nThird term: b\u00b2\n\nAdding them three terms:\n\n90,000 - ,200b +4b\u00b2 +300b -2b\u00b2 +b\u00b2\n\nCombine like terms:\n\n90,000 + (-1,200b + 00b) + (4b\u00b2 -2b\u00b2 +b\u00b2)\n\n= 90,000 - 900b + b\u00b2\n\nTherefore, the equation becomes:\n\nb*(90,000 - 900b +3b\u00b2) = 3,000,000\n\nLet me write this as:\n\n3b\u00b3 - 900b\u00b2 +90,000b -3,000,000 =0\n\nDivide both sides by 3:\n\nb\u00b3 - 300b\u00b2 +30,000b -1,000,000 =0\n\nWait, this is the same cubic equation as before when I set a = b!\n\nIndeed, the cubic is the same as when we set a = b. earlier. which had a root at b =100, and fact factored into (b -100)^3 =0. Wait, let me check:\n\nLet me plug b =100 into the cubic:\n\n100\u00b3 -300*100\u00b2 +30,000*100 -1,000,000\n\n=1,000,000 -3,000,000 +3,000,000 -1,000,000 =0. So yes, b=100 is a root.\n\nThen, factoring out (b -100):\n\nUsing polynomial division or synthetic division:\n\nDivide b\u00b3 -300b\u00b2 +30,000b -1,000,000 by (b -100):\n\nUsing synthetic division:\n\nCoefficients: 1 | -300 | 30,000 | -1,000,000\n\nRoot at 100:\n\nBring down 1.\n\nMultiply by 100: 100. Add to next300: -200\n\nMultiply -200 by 100: -20,000. Add to 30,000: 10,000\n\nMultiply 10,000 by 100: 1,000,000. Add to -1,000,000: 0.\n\nThus the cubic factors as (b -100)(b\u00b2 -200b +10,000). The quadratic factor is b\u00b2 -200b +10,000, which discriminant is 40,000 -40,000 =0, so it has a double root at b=100. Thus the equation is (b -100)^3 =0, so the only solution is b=100.\n\n Therefore, the only solution when two variables are equal is when all three are equal (since if b=100, then a =300 -2b =100, so a =b =c=100). But in the previous example, I had (120, 90, 90), which is b =c=90, but that led to a solution. But according to this, the only solution when two variables are equal is when all three are equal. Contradiction?\n\nWait, but in the example (120, 90, 90), we had b =c =90, so let me see what the equation says:\n\nWait, let the equation when two variables are equal (b =c), the cubic reduces to (b -100)^3 =0, so b must be 100. But in the example, b =90, which is't satisfy that. So there must be an error in my calculations.\n\nWait, let me reccheck the example (120, 90, 90):\n\nCompute the second equation:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\nWith a=120, b=90, c=90:\n\na\u00b2b =120\u00b2*90 =14,400*90 =1,296,000\n\na\u00b2c = same as a\u00b2b =1,296,000\n\nb\u00b2a =90\u00b2*120 =8,100*120 =972,000\n\nb\u00b2c =90\u00b2*90 =8,100*90 =729,000\n\nc\u00b2a = same as b\u00b2a =972,000\n\nc\u00b2b = same as b\u00b2c =729,000\n\nTotal sum:\n\n1,296,000 +1,296,000 =2,592,000\n\n+972,000 +729,000 = 1,701,000 \u2192 total 4,293,000\n\n+972,000 +729,000 = another 1,701,000 \u2192 total 6,000,000. So it works.\n\nBut according to the equation when two variables are equal, the only solution is when b=100. So why does this work?\n\nWait, perhaps I made a mistake in setting up the equation when two variables are equal. Let me re-derive it.\n\nLet me suppose that two variables are equal, say b =c. Then a =300 - 2b.\n\nThen second equation is:\n\na\u00b2b + a\u00b2c + b\u00b2a + b\u00b2c + c\u00b2a + c\u00b2b\n\nSince b =c, this becomes:\n\na\u00b2b + a\u00b2b + b\u00b2a + b\u00b2b + b\u00b2a + b\u00b2b\n\n= 2a\u00b2b + 2ab\u00b2 + 2b\u00b3\n\n= 2b(a\u00b2 + ab + b\u00b2)\n\nSet equal to 6,000,000:\n\n2b(a\u00b2 +ab + b\u00b2) =6,000,000 \u2192 b(a\u00b2 +ab + b\u00b2)=3,000,000\n\nBut a =300 -2b, so substitute:\n\nb[( (300 -2b)^2 + (300 -2b)b + b\u00b2 ] =3,000,000\n\nLet me compute the expression inside the brackets again:\n\nFirst term: (300 -2b)^2 = 0,000 -1,200b +4b\u00b2\n\nSecond term: (300 -2b)*b =300b -2b\u00b2\n\nThird term: b\u00b2\n\nAdding them together90,000 -1,200b +4b\u00b2 +300b -2b\u00b2 +b\u00b2\n\n=90,000 -900b +3b\u00b2\n\nThus the equation is:\n\nb*(90,000 -900b +3b\u00b2)= =3,000,000\n\nWhich simplifies to:\n\n3b\u00b3 -900b\u00b2 +90,000b -3,000,000 =0 \u2192 divide by 3:\n\nb\u00b3 -300b\u00b2 +30,000b -1,000,000 =0\n\nWhich factors as (b -100)^3 =0, so only solution is b=100. But in the example, b=90,. a solution. Contr contradiction must be an error here.\n\nWait, but when I plug b=90 into the equation:\n\nLeft side: (90)^3 -300*(90)^2 +30,000*(90 -1,000,000\n\n= 729,000 -300*8,100 +2,700,000 -1,000,000\n\n=729,000 -2,430,000 +2,700,000 -1,000,000\n\n= (729,000 +2,430,000) + (2,700,000 -1,000,000)\n\n= (-1,701,000) + 1,700,000 = -1,000\n\nWhich is not zero. So that suggests that my example (120,90,90) does not satisfy the equation? But when I computed manually, it did. Wait must be a miscalculation.\n\nWait, let me recompute the example:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n= (120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\nWait, let me compute each term step-by step:\n\nFirst term: a\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\nSecond term: a\u00b2c = same as first term:1,296,000\n\nThird term: b\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nFourth term: b\u00b2c =90\u00b2 *90 =8,100 *90 =729,000\n\nFifth term: c\u00b2a = same as third term 972,000\n\nSixth term: c\u00b2b = same as fourth term 729,000\n\nNow adding all terms:\n\nFirst two terms: 1,296,000 +1,296,000 =2,592,000\n\nThird and fourth: 72,000 +729,000 =1,701,000 \u2192 total so far 4,293,000\n\nFifth and sixth:972,000 +729,000 =1,701,000 \u2192 total 6,000,000. So it does add up.\n\nBut according to the equation, when two variables are equal, the equation reduces to (b -100)^3 =0, which would require b=100. So why is there a discrepancy?\n\nAh! Wait a second, perhaps think I made a mistake in the substitution. Let me recheck the substitution step.\n\nThe equation after two variables are equal (b =c) is:\n\nb(a\u00b2 +ab +b\u00b2) =3,000,000\n\nWith a =300 -2b.\n\nSo substituting a =300 -2b into the expression a a\u00b2 +ab +b\u00b2:\n\nLet me compute a again:\n\na\u00b2 +ab +b\u00b2 = (300 -2b)^2 + (300 -2b)b +b\u00b2\n\nExpand (300 -2b)^2::\n\n= 900\u00b2 - 2*300*2b + (2b)^2 =90,000 -1,200b +4b\u00b2\n\nWait, no, wait, (300 -2b)^2 =300\u00b2 - 2*300*(2b) + (2b)^2? Wait, wait's incorrect.\n\nWait, (a - b)^2 =a\u00b2 -2ab +b\u00b2. So (300 -2b)^2 =300\u00b2 -2*300*(2b) + (2b)^2 =90,000 -1,200b +4b\u00b2. That part is correct.\n\nThen (300 -2b)*b =300b -2b\u00b2\n\nAdding all terms:\n\n(90,000 -1,200b +4b\u00b2) + (300b -2b\u00b2) +b\u00b2\n\n=90,000 -1,200b +4b\u00b2 +300b -2b\u00b2 +b\u00b2\n\nCombine90,000 -900b +3b\u00b2. That is correct.\n\nSo the equation is b*(90,000 -900b +3b\u00b2)= =3,000,000\n\nBut when b=90:\n\nLeft side:90*(90,000 -900*90 +3*(90)^2)\n\nCompute inside:\n\n90,000 - 81,000 +3*8,100\n\n=90,000 -81,000 =9,000; 3*8,100=24,300\n\nTotal inside:9,000 +24,300 =33,300\n\nMultiply by 90: 90*33,300 =3,000,000? Let me compute:\n\n33,300 *90 = (33,000 *90) + (300 *90) =2,970,000 +27,000 =2,997,000. Which is not 3,000,000. So discrepancy is a discrepancy here here.\n\nWait, but when I computed manually, the total was 6,000,000. So there must be an error in my manual calculation.\n\nWait, let me recompute the example:\n\na=120, b=90, c=90.\n\nCompute each term:\n\na\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000\n\nc\u00b2a = same as b\u00b2a =972,000\n\nc\u00b2b = same as b\u00b2c =729,000\n\nAdding all terms:\n\n1,296,000 +1,296,000 =2,592,000\n\n+972,000 +729,000 =1,701,000 \u2192 total 4,293,000\n\n+972,000 +729,000 =1,701,000 \u2192 total 6,000,000.\n\nWait, but according to the equation, it should be 2,997,000. So there is a contradiction. Therefore I must have made a mistake in the equation setup.\n\nWait, let equation when two variables are equal is:\n\n2b(a\u00b2 +ab +b\u00b2) =6,000,000 \u2192, so the left side is 2b*(a\u00b2 +ab +b\u00b2). In the example, a=120, b=90:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300 =5,994,000, which is not 6,000,000. Wait when I computed manually, I got 6,000,000. So there must be an error in manual manual calculation.\n\nWait, let me recompute the terms:\n\na\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total for a\u00b2 terms: 2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000 \u2192 total for b\u00b2 terms: 1,701,000\n\nc\u00b2a =90\u00b2 *120 = same as b\u00b2a =972,000\n\nc\u00b2b =90\u00b2 *90 = same as b\u00b2c =729,000 \u2192 total for c\u00b2 terms: 1,701,000\n\nTotal sum:2,592,000 +1,701,000 +1,701,000 =6,000,000? Wait, 2,592,000 +1,701,000 is=4,293,000 +1,701,000 =6,000,000. But according to the equation, it should be 5,994,000. There is a discrepancy of 6,000,000 -5,994,000 =6,000.\n\nWait, so where is the mistake?\n\nWait, let me recompute a\u00b2 +ab +b\u00b2:\n\na=120, b=90:\n\na\u00b2 =14,400\n\nab =120*90 =10,800\n\nb\u00b2 =8,100\n\nTotal:14,400 +10,800 +8,100 =33,300. Correct that's correct.\n\nThen 2b*(a\u00b2 +ab +b\u00b2) =2*90*33,300 =180*33,300 =5,994,000. But manual calculation gives 6,000,000. So there is a mistake in the manual calculation.\n\nWait, let me recompute the terms again carefullya\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total a\u00b2 terms: 2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000 \u2192 total b\u00b2 terms: 1,701,000\n\nc\u00b2a =90\u00b2 *120 = same as b\u00b2a =972,000\n\nc\u00b2b =90\u00b2 *90 = same as b\u00b2c =729,000 \u2192 total c\u00b2 terms: 1,701,000\n\nTotal sum:2,592,000 +1,701,000 +1,701,000 =6,000,000.Wait according to the equation, it should be 5,994,000. There is a discrepancy of 6,000. So where must be an error in the equation setup.\n\nWait, perhaps me check the equation again.The second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b =6,000,000.\n\nWhen two variables are equal ( say b =c, then:\n\na\u00b2b +a\u00b2c =a\u00b2b2b\n\nb\u00b2a +b\u00b2c =b\u00b2a +b\u00b3b =b\u00b2(a +b\u00b3\n\nc\u00b2a +c\u00b2b =c\u00b2a +c\u00b2b = same as b\u00b2a +b\u00b2b =b\u00b2a +b\u00b3\n\nSo total equation2a\u00b2b + 2b\u00b2a + 2b\u00b3 =2a\u00b2b +2ab\u00b2 +2b\u00b3 =2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, in the example, a I computed manually the terms, I got 6,000,000, but according to the equation, it should be 5,994,000. So there must be an arithmetic error in the manual calculation.\n\nWait, let me recompute the terms:\n\nFirst=120, b=90, c=90.\n\na\u00b2b =120\u00b22 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total a\u00b2 terms: 2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =9,100 *90 =729,000 \u2192 total b\u00b2 terms: 1,701,000\n\nc\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nc\u00b2b =90\u00b2 *90 =8,100 *90 =729,000 \u2192 total c\u00b2 terms: 1,701,000\n\nTotal sum:2,592,000 +1,701,000 +1,701,000 =6,000,000\n\nWait, but according to the equation, it should be 2b*(a\u00b2 +ab +b\u00b2)= =2*90*(14,400 +10,800 +8,100)= 180*(33,300)=5,994,000. So there is a discrepancy of 6,000,000 -5,994,000 =6,000.\n\nWait, so where must be an mistakeiscalculation in the equation setup.\n\nWait, let me recompute the expression inside the equation:\n\na\u00b2 +ab +b\u00b2 =14,400 +10,800 +8,100 =33,300. Correct's correct.\n\nThen 2b*(a\u00b2 +ab +b\u00b2)= =2*90*33,300 =180*33,300.\n\nLet me compute 33,300 *180:\n\n33,300 *100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000.\n\nBut manual calculation gives 6,000,000. So there must be an error in the manual calculation.\n\nWait, let me recompute the terms again:\n\na\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total a\u00b2 terms: 2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000 \u2192 total b\u00b2 terms: 1,701,000\n\nc\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nc\u00b2b =90\u00b2 *90 =8,100 *90 =729,000 \u2192 total c\u00b2 terms: 1,701,000\n\nAdding all terms:\n\n2,592,000 +1,701,000 =4,293,000\n\n+1,701,000 =6,000,000\n\nWait, but according to the equation, it should be 5,994,000. So there is a discrepancy of 6,000. Where suggests't make sense. Unless I made a mistake in the equation setup?\n\nWait, let me check the original equation again:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b =6,000,000.\n\nWhen two variables are equal, say b =c, then:\n\na\u00b2b +a\u00b2c =a\u00b2(b +c) =a\u00b2*(2b)\n\nb\u00b2a +b\u00b2c =b\u00b2a +b\u00b2b =b\u00b2(a +b)\n\nc\u00b2a +c\u00b2b =c\u00b2a +c\u00b2b =c\u00b2(a +b) =b\u00b2(a +b)\n\nSo total:\n\n2a\u00b2b + \u00b2(a +b) +b\u00b2(a +b) =2a\u00b2b +2b\u00b2(a +b)\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3 =2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90. Sum is 120+90+90=300. Correct's correct.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\nWait, but 120 squared is 14,400, yes thatFirst two terms:14,400*90*2 =28,800*90 =2,592,000\n\nNext two terms:90\u00b2*120 +90\u00b2*90 =8,100*120 +8,100*90 =972,000 +729,000 =1,701,000\n\nLast two terms: same as the next two terms: another72,000 +729,000 =1,701,000\n\nTotal:2,592,000 +1,701,000 +1,701,000 =6,000,000\n\nBut according to the equation, it should be 2,994,000. So there must be an mistake in the equation setup.\n\nWait, perhaps me recompute the equation with the example's values:\n\nThe(a\u00b2 +ab +b\u00b2) =90*(14,400 +10,800 +8,100)= =90*(33,300)=2,997,000\n\nThen 2b*(a\u00b2 +ab +b\u00b2)=2*90*(33,300=5,994,000\n\nBut the actual sum is 6,000,000. So there is a difference of 6,000. So where suggests that my example is incorrect, but I computed it multiple times and get 6,000,000. So perhaps must be an error in the equation setup.\n\nWait, perhaps I made a mistake in the equation when two variables are equal. Let me re-derive it.\n\nLet me suppose b =c, so variables are (a, b, b.\n\n with a +2b =300.\n\nThe second equation is:\n\na\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct.\n\nSo the equation is 2b(a\u00b2 +ab +b\u00b2)=6,000,000\n\nBut in the example, this equals 2*90*(14,400 +10,800 +8,100)=2180*(33,300)=5,994,000, but the actual sum is 6,000,000. So there must be an mistake in the example.\n\nWait, perhaps I made a mistake in the example's variables. Let me check:\n\nWait, if a=120, b=90, c=90, then a120+90+90=300. That's correct.\n\nWait, let me computecompute the terms:\n\na\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total a\u00b2 terms:2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000 \u2192 total b\u00b2 terms:1,701,000\n\nc\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nc\u00b2b =9,100 *90 =729,000 \u2192 total c\u00b2 terms:1,701,000\n\nTotal sum:2,592,000 +1,701,000 +1,701,000 =6,000,000\n\nBut according to the equation, it should be 5,994,000. So there is a discrepancy of 6,000. This is very confusing.\n\nWait, perhaps I made a mistake in the equation setup. Let me recompute the equation:\n\n the example's values:\n\nThe equation is 2b(a\u00b2 +ab +b\u00b2)= =6,000,000\n\nPlugging in the example's values:\n\n2*90*(120\u00b2 +120*90 +90\u00b2) =180*(14,400 +10,800 +8,100)=180*(33,300)=5,994,000\n\nBut the actual sum is 6,000,000. So there is a mistake somewhere.\n\nWait, perhaps I made a mistake in the arithmetic:\n\nWait, 33,300 *180:\n\n33,300 * 00 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nYes.\n\nBut the manual calculation gives 6,000,000. So there must be an error in the manual calculation.\n\nWait, let me computecompute the terms:\n\n one more time:\n\na\u00b22b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total a\u00b2 terms:2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000 \u2192 total b\u00b2 terms:1,701,000\n\nc\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nc\u00b2b =8,100 *90 =729,000 \u2192 total c\u00b2 terms:1,701,000\n\nAdding all terms:\n\n2,592,000 +1,701,000 =4,293,000\n\n+1,701,000 =6,000,000\n\nWait, but according to the equation, it should be 5,994,000. So there is a discrepancy of 6,000. This is perplexing.\n\nWait, perhaps I made a mistake in the equation setup. Let me recheck the equation:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b =When b =c, this becomes:\n\na\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\nYes2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct.\n\nSo the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the example's variables. Let me check:\n\n the values again:\n\na=120, b=90, c=90.\n\nCompute a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b:\n\na\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total a\u00b2 terms:2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000 \u2192 total b\u00b2 terms:1,701,000\n\nc\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nc\u00b2b =8,100 *90 =729,000 \u2192 total c\u00b2 terms:1,701,000\n\nTotal sum:2,592,000 +1,701,000 +1,701,000 =6,000,000\n\nBut according to the equation, it should be 5,994,000. So there must be an error in the equation setup.\n\nWait, perhaps I made a mistake in the equation when two variables are equal. Let me re-check:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b.\n\nIf b =c, then:\n\na\u00b2b +a\u00b2c =a\u00b2(b +c) =a\u00b2*(2b)\n\nb\u00b2a +b\u00b2c =b\u00b2a +b\u00b2b =b\u00b2(a +b)\n\nc\u00b2a +c\u00b2b =c\u00b2a +c\u00b2b =c\u00b2(a +b) =b\u00b2(a +b)\n\nSo total:\n\n2a\u00b2b +b\u00b2(a +b) +b\u00b2(a +b) =2a\u00b2b +2b\u00b2(a +b)\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3 =2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic of calculating the example's terms. Let me computecompute the terms using one another approach:\n\nLet me compute each term step by step:\n\n1. a\u00b2b = (20\u00b2 *90 =14,400 *90 =1,296,000\n\n2. a\u00b2c =14,400 *90 =1,296,000\n\n3. b\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\n4. b\u00b2c =8,100 *90 =729,000\n\n5. c\u00b2a =8,100 *120 =972,000\n\n6. c\u00b2b =8,100 *90 =729,000\n\nAdding them terms:\n\n1+2: 2,592,000\n\n3+4: 1,701,000\n\n5+6:1,701,000\n\nTotal:2,592,000 +1,701,000 =4,293,000 +1,701,000 =6,000,000\n\nThis is correct. So the equation must be wrong.Wait, but according to the equation, it should be 5,994,000. So there must be a mistake in the equation setup.\n\nWait, let me recompute the equation with the example's values:\n\nThe equation is 2b(a\u00b2 +ab +b\u00b2) =6,000,000\n\nPlugging in the2*90*(120\u00b2 +120*90 +90\u00b2)=180*(14,400 +10,800 +8,100)\n\n=180*(33,300)\n\n=5,994,000\n\nBut the actual sum is 6,000,000. So there is a discrepancy of 6,000. This suggests that there is an error in the equation setup.\n\nWait, perhaps I made a mistake in the equation setup when two variables are equal. Let me re-derive it againLet me recompute the expression:\n\nWhen two =c,The second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +222b\u00b2a + 2b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(109,800) +2*(729,000)\n\nWait, wait,   think I made a mistake here:\n\nWait, 90\u00b2*120 is 8,100*120 =972,000, not 109,800. So:\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 = =2,592,000 +1,944,000 =4,536,000 +1,458,000 =6,000,000\n\nWait according to the equation:\n\n2b(a\u00b2 +ab +b\u00b2)= =2*90*(14,400 +10,800 +8,100)=180*(33,300)=5,994,000\n\nWait there is a discrepancy of 6,000,000 -5,994,000 =6,000.\n\nWait, but where terms are:\n\n2*(a\u00b2b) +2*(b\u00b2a) +2*(b\u00b3)\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3 = =2b(a\u00b2 +ab +b\u00b2)\n\nYes in the example:\n\n2a\u00b2b =2*14,400*90 =2*1,296,000 =2,592,000\n\n2ab\u00b2 =2*120*8,100 =2*972,000 =1,944,000\n\n2b\u00b3 =2*(729,000 =1,458,000\n\nTotal:2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nBut according to the equation,2b(a\u00b2 +ab +b\u00b2)=2b90*(14,400 +10,800 +8,100)=2*90*(33,300)=5,994,000\n\nWait, but 2b*(\u00b2 +ab +b\u00b2)= is 2*90*(14,400 +10,800 +8,100)= 180*(33,300)=5,994,000\n\nBut the actual terms are 6,000,000. So there must be an mistake in the equation setup.\n\nWait, perhaps's possible that I made a mistake in the equation setup. Let me recompute the expression inside the equation:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nYes, that's correct.\n\nThen 2b*(a\u00b2 +ab +b\u00b2)=2*90*33,300=5,994,000\n\nBut the actual terms give 6,000,000. So there is a mistake somewhere the equation setup.\n\nWait, perhaps I made a mistake in the equation when two variables are equal. Let me re-derive it againThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c,= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b ++a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nBut according to the equation, it should be 5,994,000. So there is a discrepancy of 6,000. This is't possible. unless there is a calculation error.\n\nWait, let me compute 20\u00b2*90:\n\n120 squared is 14,400. 14,400 *90 =1,296,000. That's correct.\n\n90 squared is 8,100. 8,100 *120 =972,000. Correct.\n\n8,100 *90 =729,000. Correct.\n\nSo the terms are correct. So the equation must be wrong.\n\nWait, perhaps I made a mistake in the equation setup. Let me recheck:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c, this is:\n\na\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me computecompute 2b(a\u00b2 +ab +b\u00b2):\n\nWith b=90, a=120:\n\na\u00b2 +ab +b\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply\n\nOkay, I think I need to take a step back. here. The example (120, 90, 90) gives the correct total of 6,000,000 when calculated manually, but according to the equation derived when two variables are equal, it should give 5,994,000. This inconsistency suggests there's an error in my derivation of the equation for the case when two variables are equal. However, after multiple checks, the derivation seems correct. This contradiction't make sense. Perhaps I made a mistake in the example's variables?\n\nWait, let me try another example where two variables are equal. Let's take b =c =100, then a =300 -200 =100. So the triple is (100,100,100). which should work. Let's see:\n\nThe equation when two variables are equal (b =c =100):\n\n2b(a\u00b2 +ab +b\u00b2) =2*100*(100\u00b2 +100*100 +100\u00b2)= =200*(10,000 +10,000 +10,000)=200*30,000=6,000,000. Which matches the required value. So that works.\n\nAnother example: let me take b =c =50, so a =300 -100 =200. So the triple is (200,50,50). Let's compute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=200\u00b2*50 +200\u00b2*50 +50\u00b2*200 +50\u00b2*50 +50\u00b2*200 +50\u00b2*50\n\n=2*(40,000*50) +2*(2,500*200) +2*(2,500*50)\n\n=2*(2,000,000 +2*500,000 +2*125,000\n\n=4,000,000 +1,000,000 +250,000 =5,250,000\n\nAccording to the equation:\n\n2b(a\u00b2 +ab +b\u00b2) =2*50*(200\u00b2 +200*50 +50\u00b2)\n\n=100*(40,000 +10,000 +2,500)=100*(52,500)=5,250,000. Which matches.\n\nSo the equation works for this case. But in the case of b=90, the gives a discrepancy. So why?\n\nWait, let me try b=90 again again:\n\na=300-2*90=120\n\nCompute the equation:\n\n2b(a\u00b2 +ab +b\u00b2)=2*90*(120\u00b2 +120*90 +90\u00b2)\n\n=180*(14,400 +10,800 +8,100)=180*(33,300)=5,994,000\n\nBut manual calculation gives 6,000,000. So there must be an error in the manual calculation.\n\nWait, let me recompute the terms for (120,90,90):\n\na\u00b2b =120\u00b2 *90 =14,400 *90 =1,296,000\n\na\u00b2c =14,400 *90 =1,296,000 \u2192 total a\u00b2 terms:2,592,000\n\nb\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nb\u00b2c =8,100 *90 =729,000 \u2192 total b\u00b2 terms:1,701,000\n\nc\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\nc\u00b2b =8,100 *90 =729,000 \u2192 total c\u00b2 terms:1,701,000\n\nTotal sum:2,592,000 +1,701,000 +1,701,000 =6,000,000\n\nBut according to the equation, it should be 5,994,000. So is a contradiction. I must have made a mistake in the equation setup.\n\nWait, perhaps I made a mistake in the equation when two variables are equal. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\nWait, wait, 90\u00b2*120 is 8,100*120 =972,000, not 972,000. So:\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nBut according to the equation:\n\n2b(a\u00b2 +ab +b\u00b2)=2*90*(14,400 +10,800 +8,100)=180*(33,300)=5,994,000\n\nThere is a discrepancy of 6,000. This suggests't make sense. I must be missing something.\n\nWait, perhaps I made a mistake in the equation setup. Let me recompute the expression inside when two variables are equal:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\nLet3,300 *180:\n\n33,300 *100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is be an error in the manual calculation.\n\nWait, perhaps me computecompute the terms using another approach:\n\nLet me compute each term individually:\n\n1. a\u00b2b =120\u00b22 *90 =14,400 *90 =1,296,000\n\n2. a\u00b2c =14,400 *90 =1,296,000\n\n3. b\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\n4. b\u00b2c =8,100 *90 =729,000\n\n5. c\u00b2a =90\u00b2 *120 =8,100 *120 =972,000\n\n6. c\u00b2b =8,100 *90 =729,000\n\nAdding all terms:\n\n1+2 =2,592,000\n\n3+4 =1,701,000\n\n5+6 =1,701,000\n\nTotal:2,592,000 +1,701,000 =4,293,000 +1,701,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me try to compute the equation with the example's values using using the original equation:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b =6,000,000\n\nWith a=120, b=90, c=90:\n\n120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nYes according to the equation when two variables are equal:\n\n2b(a\u00b2 +ab +b\u00b2)=2*90*(14,400 +10,800 +8,100)=5,994,000\n\nThis suggests't add up. There must be a mistake in the equation setup.\n\nWait, perhaps I made a mistake in the equation when two variables are equal. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is't possible. unless there is an calculation error.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nYes is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me try to compute the equation with the example's values using the original equation:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b =6,000,000\n\nWith a=120, b=90, c=90:\n\n120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nYes is correct. So the equation must two variables are equal must be wrong. But I can't see why.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is perplexing.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me try to compute the equation with the example's values using the original equation:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b =6,000,000\n\nWith a=120, b=90, c=90:\n\n120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation when two variables are equal must be wrong. But I can't see why.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is perplex.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me try to compute the equation with the example's values using the original equation:\n\nThe second equation is a\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b =6,000,000\n\nWith a=120, b=90, c=90:\n\n120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation when two variables are equal must be wrong. But I can't see why.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it again:\n\nThe second equation is:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\nWhen b =c:\n\n= a\u00b2b +a\u00b2b +b\u00b2a +b\u00b2b +b\u00b2a +b\u00b2b\n\n=2a\u00b2b +2ab\u00b2 +2b\u00b3\n\n=2b(a\u00b2 +ab +b\u00b2)\n\nYes, that's correct. So the equation is correct. So why the discrepancy?\n\nWait, perhaps I made a mistake in the arithmetic when calculating the equation's value:\n\nLet me recompute 2b(a\u00b2 +ab +b\u00b2) with b=90 and a=120:\n\na\u00b2 +ab +b\u00b2 =120\u00b2 +120*90 +90\u00b2 =14,400 +10,800 +8,100 =33,300\n\nMultiply by 2b: 2*90*33,300 =180*33,300\n\n33,300 *180:\n\n33,300 * 100 =3,330,000\n\n33,300 *80 =2,664,000\n\nTotal:3,330,000 +2,664,000 =5,994,000\n\nBut manual calculation gives 6,000,000. So there is a discrepancy of 6,000. This is impossible.\n\nWait, perhaps I made a mistake in the example's values. Let me check the values again:\n\na=120, b=90, c=90.\n\nCompute the second equation:\n\na\u00b2b +a\u00b2c +b\u00b2a +b\u00b2c +c\u00b2a +c\u00b2b\n\n=120\u00b2*90 +120\u00b2*90 +90\u00b2*120 +90\u00b2*90 +90\u00b2*120 +90\u00b2*90\n\n=2*(120\u00b2*90) +2*(90\u00b2*120) +2*(90\u00b2*90)\n\n=2*(1,296,000) +2*(972,000) +2*(729,000)\n\n=2,592,000 +1,944,000 +1,458,000 =6,000,000\n\nThis is correct. So the equation must be wrong. But I can't see where.\n\nWait, perhaps I made a mistake in the equation setup. Let me re-derive it"
+    print(evaluator_map["aime2024"].llm_judge(response, "601", "", False))

rosetta/utils/registry.py

@@ -0,0 +1,279 @@
+"""
+Unified registry utilities and simple JSON-based save/load helpers.
+
+This module provides:
+- create_registry: factory to create (registry dict, register decorator, get_class)
+- capture_init_args: decorator to record __init__ kwargs on instances as _init_args
+- save_object / load_object: serialize/deserialize object configs via registry
+"""
+
+from __future__ import annotations
+
+import inspect
+import json
+from typing import Dict, Type, Callable, Optional, Tuple, TypeVar, Any
+import torch
+
+T = TypeVar("T")
+
+
+def create_registry(
+    registry_name: str,
+    case_insensitive: bool = False,
+) -> Tuple[Dict[str, Type[T]], Callable[..., Type[T]], Callable[[str], Type[T]]]:
+    """
+    Create a registry system with register and get functions.
+
+    Args:
+        registry_name: Name used in error messages (e.g., "projector")
+        case_insensitive: Whether to store lowercase versions of names
+
+    Returns:
+        (registry_dict, register_function, get_function)
+    """
+
+    registry: Dict[str, Type[T]] = {}
+
+    def register(cls_or_name=None, name: Optional[str] = None):
+        """Register a class in the registry. Supports multiple usage patterns.
+
+        Usage:
+            @register
+            class Foo: ...
+
+            @register("foo")
+            class Foo: ...
+
+            @register(name="foo")
+            class Foo: ...
+        """
+
+        def _register(c: Type[T]) -> Type[T]:
+            # Determine the name to use
+            if isinstance(cls_or_name, str):
+                class_name = cls_or_name
+            elif name is not None:
+                class_name = name
+            else:
+                class_name = c.__name__
+
+            registry[class_name] = c
+            if case_insensitive:
+                registry[class_name.lower()] = c
+            return c
+
+        if cls_or_name is not None and not isinstance(cls_or_name, str):
+            # Called as @register or register(cls)
+            return _register(cls_or_name)
+        else:
+            # Called as @register("name") or @register(name="name")
+            return _register
+
+    def get_class(name: str) -> Type[T]:
+        """Get class by name from registry."""
+        if name not in registry:
+            # Build readable available list without duplicates when case_insensitive
+            seen = set()
+            available = []
+            for k in registry.keys():
+                if k.lower() in seen:
+                    continue
+                seen.add(k.lower())
+                available.append(k)
+            raise ValueError(
+                f"Unknown {registry_name} class: {name}. Available: {available}"
+            )
+        return registry[name]
+
+    return registry, register, get_class
+
+
+def capture_init_args(cls):
+    """
+    Decorator to capture initialization arguments of a class.
+
+    Stores the mapping of the constructor's parameters to the values supplied
+    at instantiation time into `self._init_args` for later serialization.
+    """
+    original_init = cls.__init__
+
+    def new_init(self, *args, **kwargs):
+        # Store all initialization arguments
+        init_args: Dict[str, Any] = {}
+
+        # Get parameter names from the original __init__ method
+        sig = inspect.signature(original_init)
+        param_names = list(sig.parameters.keys())[1:]  # Skip 'self'
+
+        # Map positional args to parameter names
+        for i, arg in enumerate(args):
+            if i < len(param_names):
+                init_args[param_names[i]] = arg
+
+        # Add keyword args
+        init_args.update(kwargs)
+
+        self._init_args = init_args
+
+        # Call the original __init__
+        original_init(self, *args, **kwargs)
+
+    cls.__init__ = new_init
+    return cls
+
+
+# -------------------------
+# Serialization utilities
+# -------------------------
+
+def _encode_value(value: Any) -> Any:
+    """Best-effort JSON encoding for common ML types."""
+    # Primitives and None
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+
+    # Tuples -> lists
+    if isinstance(value, tuple):
+        return [
+            _encode_value(v) for v in value
+        ]
+
+    # Lists
+    if isinstance(value, list):
+        return [
+            _encode_value(v) for v in value
+        ]
+
+    # Dicts
+    if isinstance(value, dict):
+        return {k: _encode_value(v) for k, v in value.items()}
+
+    # torch-specific types
+    if torch is not None:
+        # torch.dtype
+        if isinstance(value, type(getattr(torch, "float32", object))):
+            # Guard: torch.dtype is not a class; rely on str(value) format
+            s = str(value)
+            if s.startswith("torch."):
+                return {"__type__": "torch.dtype", "value": s.split(".")[-1]}
+
+        # torch.device
+        if isinstance(value, getattr(torch, "device", ())):
+            return {"__type__": "torch.device", "value": str(value)}
+
+    # Fallback to string representation
+    return {"__type__": "str", "value": str(value)}
+
+
+def _decode_value(value: Any) -> Any:
+    """Decode values produced by _encode_value, recursively for containers."""
+    # Lists: decode each element
+    if isinstance(value, list):
+        return [_decode_value(v) for v in value]
+
+    # Dicts: either a typed-marker dict or a regular mapping that needs recursive decoding
+    if isinstance(value, dict):
+        if "__type__" in value:
+            t = value.get("__type__")
+            v = value.get("value")
+
+            if t == "torch.dtype" and torch is not None:
+                dtype = getattr(torch, str(v), None)
+                if dtype is None:
+                    raise ValueError(f"Unknown torch.dtype: {v}")
+                return dtype
+
+            if t == "torch.device" and torch is not None:
+                return torch.device(v)
+
+            if t == "str":
+                return str(v)
+
+            # Unknown type marker; return raw as-is
+            return value
+
+        # Regular dict: decode values recursively
+        return {k: _decode_value(v) for k, v in value.items()}
+
+    # Primitives and anything else: return as-is
+    return value
+
+
+def save_object(obj: Any, file_path: str) -> None:
+    """
+    Save an object's construction config to a JSON file.
+
+    The object is expected to have been decorated with capture_init_args,
+    so that `obj._init_args` exists.
+    """
+    class_name = obj.__class__.__name__
+    init_args = getattr(obj, "_init_args", {})
+
+    serializable_args = _encode_value(init_args)
+    payload = {
+        "class": class_name,
+        "init_args": serializable_args,
+    }
+
+    with open(file_path, "w", encoding="utf-8") as f:
+        json.dump(payload, f, indent=2)
+
+
+def load_object(
+    file_path: str,
+    get_class_fn: Callable[[str], Type[T]],
+    override_args: Optional[Dict[str, Any]] = None,
+) -> T:
+    """
+    Load an object from a JSON config file previously saved by save_object.
+
+    Args:
+        file_path: Path to JSON file
+        get_class_fn: Function to resolve class names from registry
+        override_args: Optional dict to override stored init args
+
+    Returns:
+        Instantiated object of type T
+    """
+    with open(file_path, "r", encoding="utf-8") as f:
+        payload = json.load(f)
+
+    class_name = payload["class"]
+    encoded_args = payload.get("init_args", {})
+    init_args = _decode_value(encoded_args)
+
+    if override_args:
+        init_args.update(override_args)
+
+    cls = get_class_fn(class_name)
+    return cls(**init_args)
+
+
+def dumps_object_config(obj: Any) -> str:
+    """Return a JSON string with the object's class and init args."""
+    class_name = obj.__class__.__name__
+    init_args = getattr(obj, "_init_args", {})
+    serializable_args = _encode_value(init_args)
+    return json.dumps({"class": class_name, "init_args": serializable_args}, indent=2)
+
+
+def loads_object_config(
+    s: str,
+    get_class_fn: Callable[[str], Type[T]],
+    override_args: Optional[Dict[str, Any]] = None,
+) -> T:
+    """Instantiate an object from a JSON string produced by dumps_object_config."""
+    payload = json.loads(s)
+    class_name = payload["class"]
+    encoded_args = payload.get("init_args", {})
+    init_args = _decode_value(encoded_args)
+    if override_args:
+        init_args.update(override_args)
+    cls = get_class_fn(class_name)
+    return cls(**init_args)
+
+
+# Model Registry System (case-insensitive for backward compatibility)
+PROJECTOR_REGISTRY, register_model, get_projector_class = create_registry(
+    "projector", case_insensitive=True
+)