API Reference for model-training¶

def get_training_config(
    task_type: str,
    rank: int = 64,
    epochs: int = 3,
    learning_rate: float = 2e-4,
) -> dict[str, Any]:
    """Return a training configuration dict with hyperparameters.

    Generates a configuration appropriate for the given task type,
    with sensible defaults for QLoRA fine-tuning.

    Args:
        task_type: Task category (e.g. 'bug-fix', 'feature-impl').
        rank: LoRA rank for the adapter.
        epochs: Number of training epochs.
        learning_rate: Learning rate for the optimizer.

    Returns:
        A dict containing all training hyperparameters.

    Example:
        >>> config = get_training_config("bug-fix", rank=64, epochs=3)
        >>> config["task_type"]
        'bug-fix'
    """
    return {
        "task_type": task_type,
        "rank": rank,
        "alpha": 2 * rank,
        "epochs": epochs,
        "learning_rate": learning_rate,
        "warmup_ratio": 0.03,
        "lr_scheduler_type": "cosine",
        "bf16": True,
        "per_device_train_batch_size": 1,
        "gradient_accumulation_steps": 4,
        "save_strategy": "no",
        "logging_steps": 1,
        "report_to": "none",
        "eval_strategy": "no",
        "target_modules": ["q_proj", "v_proj"],
        "dropout": 0.1,
    }

def validate_config(config: dict[str, Any]) -> bool:
    """Validate training configuration fields and value ranges.

    Checks that all required fields are present and their values
    fall within acceptable ranges.

    Args:
        config: A training configuration dict to validate.

    Returns:
        True if the configuration is valid.

    Raises:
        ValueError: If required keys are missing or values are out of range.

    Example:
        >>> valid = validate_config({"task_type": "bug-fix", "rank": 64,
        ...                         "epochs": 3, "learning_rate": 2e-4})
        >>> valid
        True
    """
    missing = _REQUIRED_KEYS - set(config.keys())
    if missing:
        raise ValueError(f"Missing required config keys: {sorted(missing)}")

    rank = config["rank"]
    if not (isinstance(rank, int) and rank > 0 and rank <= 256):
        raise ValueError(f"rank must be an integer in range (0, 256], got {rank!r}")

    epochs = config["epochs"]
    if not (isinstance(epochs, int) and epochs > 0 and epochs <= 100):
        raise ValueError(f"epochs must be an integer in range (0, 100], got {epochs!r}")

    lr = config["learning_rate"]
    if not (isinstance(lr, float) and lr > 0.0 and lr < 1.0):
        raise ValueError(f"learning_rate must be a float in range (0, 1), got {lr!r}")

    return True

def get_d2l_qwen3_config() -> dict[str, Any]:
    """Return Qwen3-Coder-Next architecture dimensions without loading model weights.

    Uses Qwen3NextConfig defaults which exactly match Qwen3-Coder-Next specs:
    - hidden_size: 2048
    - num_hidden_layers: 48 (12 full_attention + 36 linear_attention)
    - num_attention_heads: 16 (Q heads), num_key_value_heads: 2 (GQA KV)
    - head_dim: 256
    - full_attention layer indices: [3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43, 47]
    - vocab_size: 151936
    - model_type: "qwen3_next"

    Returns:
        Dict with keys: hidden_size, num_hidden_layers, num_attention_heads,
        num_key_value_heads, head_dim, attention_layer_indices, vocab_size,
        model_type.
    """
    from transformers import Qwen3NextConfig  # noqa: PLC0415

    cfg = Qwen3NextConfig()
    layer_types: list[str] = cfg.layer_types or []
    attention_layer_indices = [
        i for i, t in enumerate(layer_types) if t == "full_attention"
    ]
    return {
        "hidden_size": cfg.hidden_size,
        "num_hidden_layers": cfg.num_hidden_layers,
        "num_attention_heads": cfg.num_attention_heads,
        "num_key_value_heads": cfg.num_key_value_heads,
        "head_dim": cfg.head_dim,
        "attention_layer_indices": attention_layer_indices,
        "vocab_size": cfg.vocab_size,
        "model_type": cfg.model_type,
    }

def build_qwen3_hypernet_config(
    lora_r: int = 8,
    target_modules: list[str] | None = None,
    aggregator_config: Any = None,
) -> Any:
    """Construct HypernetConfig targeting Qwen3-Coder-Next attention layers.

    Discovers full_attention layer indices dynamically from Qwen3NextConfig.layer_types.
    Result has exactly 12 layer indices matching the Qwen3-Coder-Next architecture.

    Phase 26 probe cache integration: if a probe cache exists for
    QWEN3_NEXT_CANONICAL_NAME, uses real per-projection in/out dimensions for
    feature_sizes. Falls back to hidden_size placeholder when no cache is found
    (e.g., in CI where the model has not been probed).

    Args:
        lora_r: LoRA rank for the adapter. Defaults to 8.
        target_modules: LoRA target module names. Defaults to ["q_proj", "v_proj"].
        aggregator_config: Perceiver aggregator config from a Sakana checkpoint.
            If None (default / Phase 25 CI), HypernetConfig is built with
            aggregator_config=None as placeholder. Phase 29 populates this via
            get_aggregator_config() with a loaded model.

    Returns:
        HypernetConfig with layer_indices set to the 12 full_attention indices
        and base_hidden_size=2048.
    """
    from ctx_to_lora.modeling.hypernet import HypernetConfig  # noqa: PLC0415
    from peft import LoraConfig  # noqa: PLC0415
    from transformers import Qwen3NextConfig  # noqa: PLC0415

    if target_modules is None:
        target_modules = ["q_proj", "v_proj"]

    cfg = Qwen3NextConfig()
    layer_types: list[str] = cfg.layer_types or []
    layer_indices = [
        i for i, t in enumerate(layer_types) if t == "full_attention"
    ]  # Always [3, 7, 11, 15, 19, 23, 27, 31, 35, 39, 43, 47]

    lora_config = LoraConfig(
        r=lora_r,
        lora_alpha=lora_r * 2,
        target_modules=target_modules,
        lora_dropout=0.0,
        bias="none",
        task_type="CAUSAL_LM",
    )

    from model_training.d2l_probe import (  # noqa: PLC0415
        QWEN3_NEXT_CANONICAL_NAME,
        load_probe_cache,
    )

    cache = load_probe_cache(QWEN3_NEXT_CANONICAL_NAME)
    if cache is not None:
        in_sizes = {mod: cache["feature_sizes"][mod]["in"] for mod in target_modules}
        out_sizes = {mod: cache["feature_sizes"][mod]["out"] for mod in target_modules}
        feature_sizes: tuple[dict[str, int], dict[str, int]] = (in_sizes, out_sizes)
        logger.info("Using probe cache feature_sizes for %s", QWEN3_NEXT_CANONICAL_NAME)
    else:
        hidden: int = cfg.hidden_size or 2048
        _placeholder: dict[str, int] = dict.fromkeys(target_modules, hidden)
        feature_sizes = (_placeholder, dict.fromkeys(target_modules, hidden))
        logger.warning(
            "No probe cache for '%s' — using hidden_size=%d as placeholder. "
            "Run probe_model() and save_probe_cache() to set real dimensions.",
            QWEN3_NEXT_CANONICAL_NAME,
            cfg.hidden_size,
        )

    return HypernetConfig(
        latent_size=512,
        use_light_weight_lora=False,
        light_weight_latent_size=128,
        per_rank_gen=False,
        use_per_rank_bias=False,
        use_bias=True,
        per_layer_processing=False,
        use_token_mixing=False,
        num_pre_head_layers=1,
        dropout_rate=0.0,
        lora_config=lora_config,
        extra_modules=None,
        base_hidden_size=cfg.hidden_size,
        layer_indices=layer_indices,
        feature_sizes=feature_sizes,
        aggregator_config=aggregator_config,
    )

def save_jsonl(records: list[dict[str, Any]], path: str | Path) -> None:
    """Persist a list of dicts to a JSONL file.

    Args:
        records: List of JSON-serializable dicts to save.
        path: File path for output. Parent directories are created if needed.
    """
    out = Path(path)
    out.parent.mkdir(parents=True, exist_ok=True)
    with out.open("w", encoding="utf-8") as f:
        for record in records:
            f.write(json.dumps(record, ensure_ascii=False) + "\n")

def load_jsonl(path: str | Path) -> list[dict[str, Any]]:
    """Load records from a JSONL file.

    Args:
        path: Path to the JSONL file.

    Returns:
        List of dicts, one per non-empty line.
    """
    src = Path(path)
    return [
        json.loads(line)
        for line in src.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]

def format_for_distillation(trajectory: dict[str, Any]) -> list[dict[str, str]]:
    """Convert a trajectory to distillation records with activation/teacher split.

    Each record has:
    - activation_text: trajectory context + task description (NO answer tokens)
    - teacher_text: trajectory context + task description + answer
    - task_id: identifier for train/test splitting

    Only successful trajectories (outcome == 'success') produce records.

    Args:
        trajectory: Trajectory dict with task_id/session_id, task_description,
            steps, and outcome fields.

    Returns:
        List of distillation record dicts, or empty list if no successful outcome.
    """
    if trajectory.get("outcome") != "success":
        return []

    steps: list[dict[str, Any]] = trajectory.get("steps", [])
    task_description: str = trajectory.get("task_description", "")
    task_id: str = trajectory.get("task_id") or trajectory.get("session_id", "")

    # Build trajectory context from step descriptions (no answer tokens)
    trajectory_parts: list[str] = []
    for step in steps:
        desc = step.get("description", "")
        if desc:
            trajectory_parts.append(desc)

    trajectory_text = "\n".join(trajectory_parts)

    # Build activation_text (context only, no answer)
    activation_base = f"{trajectory_text}\n{task_description}".strip()

    records: list[dict[str, str]] = []
    for step in steps:
        if not step.get("tests_passed"):
            continue

        # Extract answer from canonical_solution or generated_code
        answer: str = step.get("canonical_solution") or step.get("generated_code", "")
        if not answer:
            continue

        records.append(
            {
                "task_id": task_id,
                "activation_text": activation_base,
                "teacher_text": f"{activation_base}\n{answer}",
            }
        )

    return records

def normalize_mined_trajectory(mined: dict[str, Any]) -> dict[str, Any]:
    """Convert a GitHub-mined trajectory dict into distillation-ready format.

    Maps the mining pipeline's output (PR/issue metadata with commit and review
    steps) into the schema expected by ``format_for_distillation``.

    Outcome mapping:
        - ``pr_*`` task_ids: ``"merged"`` -> ``"success"``, else ``"failure"``
        - ``issue_*`` task_ids: ``"closed"`` -> ``"success"``, else ``"failure"``
        - Fallback: ``"merged"`` or ``"closed"`` -> ``"success"``, else ``"failure"``

    Step mapping:
        - Commit steps get ``[Commit]`` prefix in description and their content
          as ``generated_code``.
        - Review steps get ``[Review]`` prefix with content inlined into
          description and empty ``generated_code``.
        - Only the **last** commit step receives ``tests_passed=True`` and
          ``canonical_solution`` (and only when the overall outcome is success).

    Args:
        mined: Dict from the GitHub mining pipeline with keys ``task_id``,
            ``task_description``, ``outcome``, and ``steps`` (list of dicts
            with ``type``, ``description``, and ``content`` fields).

    Returns:
        Trajectory dict ready for ``format_for_distillation``.
    """
    task_id: str = mined.get("task_id", "")
    raw_outcome: str = mined.get("outcome", "")
    raw_steps: list[dict[str, Any]] = mined.get("steps", [])

    # --- Determine normalized outcome ---
    if task_id.startswith("pr_"):
        normalized_outcome = "success" if raw_outcome == "merged" else "failure"
    elif task_id.startswith("issue_"):
        normalized_outcome = "success" if raw_outcome == "closed" else "failure"
    else:
        normalized_outcome = (
            "success" if raw_outcome in ("merged", "closed") else "failure"
        )

    # --- Identify commit steps to find the last one ---
    commit_steps = [s for s in raw_steps if s.get("type") == "commit"]

    # --- Normalize steps ---
    normalized_steps: list[dict[str, Any]] = []
    for step in raw_steps:
        step_type = step.get("type", "")
        step_description = step.get("description", "")
        step_content = step.get("content", "")

        if step_type == "commit":
            is_last_commit = commit_steps and step is commit_steps[-1]
            tests_passed = is_last_commit and normalized_outcome == "success"
            entry: dict[str, Any] = {
                "description": f"[Commit] {step_description}",
                "generated_code": step_content,
                "tests_passed": tests_passed,
            }
            if is_last_commit and normalized_outcome == "success":
                entry["canonical_solution"] = step_content
            normalized_steps.append(entry)
        elif step_type == "review":
            normalized_steps.append(
                {
                    "description": f"[Review] {step_content}",
                    "generated_code": "",
                    "tests_passed": False,
                }
            )

    return {
        "task_id": task_id,
        "session_id": task_id,
        "task_description": mined.get("task_description", ""),
        "steps": normalized_steps,
        "outcome": normalized_outcome,
    }

def generate_needle_dataset(n: int = 20) -> list[dict[str, str]]:
    """Generate needle-in-haystack records for CI smoke testing.

    Records are deterministic (no randomness, no LLM). Each record contains
    a code fact embedded in a function/class context with a query and answer.

    Args:
        n: Number of records to generate. Cycles through templates if n exceeds
           the number of available templates.

    Returns:
        List of n record dicts with activation_text, teacher_text, and task_id.
    """
    records: list[dict[str, str]] = []
    n_templates = len(_NEEDLE_TEMPLATES)
    n_slots = len(_SLOT_VALUES)

    for i in range(n):
        template = _NEEDLE_TEMPLATES[i % n_templates]
        slots = _SLOT_VALUES[i % n_slots]

        try:
            trajectory = template["trajectory_template"].format(**slots)
            query = template["query_template"].format(**slots)
            answer = template["answer_template"].format(**slots)
        except KeyError:
            # If template has a slot not in _SLOT_VALUES, use a safe fallback
            trajectory = (
                f"# code fact {i}\ndef func_{i}(x: int) -> int:\n    return {i}"
            )
            query = f"What does func_{i} return?"
            answer = str(i)

        activation_text = f"{trajectory}\n\nQ: {query}"
        teacher_text = f"{trajectory}\n\nQ: {query}\nA: {answer}"

        records.append(
            {
                "task_id": f"needle_{i}",
                "activation_text": activation_text,
                "teacher_text": teacher_text,
            }
        )

    return records

def generate_trajectory_dataset(
    source: str = "humaneval",
    max_tasks: int | None = None,
) -> list[dict[str, str]]:
    """Generate trajectory records from a coding task dataset.

    Each record has activation_text (prompt only, no solution) and teacher_text
    (prompt + canonical solution), making it ready for KL-divergence distillation.

    Args:
        source: Dataset source identifier. Currently only "humaneval" is supported.
        max_tasks: Maximum number of tasks to process. If None, all tasks are used.

    Returns:
        List of trajectory record dicts with task_id, activation_text, and
        teacher_text fields.

    Raises:
        ValueError: If an unsupported source is specified.
    """
    from datasets import load_dataset  # noqa: PLC0415

    if source == "humaneval":
        dataset = load_dataset("openai_humaneval", split="test", trust_remote_code=True)
    else:
        raise ValueError(
            f"Unsupported source: {source!r}. Only 'humaneval' is supported."
        )

    records: list[dict[str, str]] = []
    tasks = list(dataset) if max_tasks is None else list(dataset)[:max_tasks]
    for task in tasks:
        task_id: str = task["task_id"]
        prompt: str = task["prompt"]
        canonical_solution: str = task["canonical_solution"]

        # activation_text: the function signature + docstring (no solution)
        activation_text = prompt.rstrip()
        # teacher_text: prompt + solution (the full implementation)
        teacher_text = prompt.rstrip() + canonical_solution

        records.append(
            {
                "task_id": task_id,
                "activation_text": activation_text,
                "teacher_text": teacher_text,
            }
        )

    logger.info("Generated %d trajectory records from %r", len(records), source)
    return records

def augment_trajectories(
    trajectories: list[dict[str, Any]],
    n_variants: int = 3,
    model: str = "qwen2.5-coder:1.5b",
    ollama_base_url: str | None = None,
) -> list[dict[str, Any]]:
    """Produce LLM-augmented variants of trajectory records.

    For each input trajectory, generates up to n_variants augmented records
    using an Ollama LLM. Augmented records always inherit the source task_id
    to preserve split integrity when mixed with originals.

    Augmentation strategies (up to n_variants selected in order):
    1. Paraphrase the task description
    2. Reorder/drop steps in the trajectory
    3. Rename variables throughout the trajectory

    Args:
        trajectories: List of trajectory record dicts with task_id,
            activation_text, and teacher_text fields.
        n_variants: Number of augmented variants to produce per trajectory.
            Maximum 3 (one per augmentation strategy).
        model: Ollama model identifier to use for augmentation.
        ollama_base_url: Ollama base URL. Defaults to "http://localhost:11434".

    Returns:
        List of augmented trajectory dicts. Each record has the same task_id
        as its source trajectory, with LLM-generated activation_text and
        teacher_text.
    """
    import asyncio  # noqa: PLC0415

    from inference.ollama_provider import OllamaProvider  # noqa: PLC0415

    base_url = ollama_base_url or "http://localhost:11434"
    provider = OllamaProvider(base_url=base_url)

    augmentation_prompts = [
        (
            "Paraphrase the following coding task description and trajectory, "
            "keeping the same logic but using different wording:\n\n{text}"
        ),
        (
            "Rewrite the following coding trajectory by reordering and slightly "
            "dropping some intermediate steps, while preserving the overall "
            "outcome:\n\n{text}"
        ),
        (
            "Rewrite the following code trajectory by renaming variables and "
            "function parameters to different names, keeping the same logic:"
            "\n\n{text}"
        ),
    ]

    async def _augment_one(
        trajectory: dict[str, Any],
        prov: Any,
        n: int,
        mdl: str,
    ) -> list[dict[str, Any]]:
        task_id: str = trajectory["task_id"]
        activation_text: str = trajectory.get("activation_text", "")
        teacher_text: str = trajectory.get("teacher_text", "")

        results: list[dict[str, Any]] = []
        strategies = augmentation_prompts[:n]
        for strategy_prompt in strategies:
            augment_prompt = strategy_prompt.format(text=teacher_text)
            result = await prov.generate(augment_prompt, mdl)
            augmented_teacher = result.text

            # Build augmented activation by stripping the solution portion
            # Use teacher text length ratio to approximate activation
            activation_prompt = strategy_prompt.format(text=activation_text)
            act_result = await prov.generate(activation_prompt, mdl)
            augmented_activation = act_result.text

            results.append(
                {
                    "task_id": task_id,  # CRITICAL: inherit source task_id
                    "activation_text": augmented_activation,
                    "teacher_text": augmented_teacher,
                }
            )
        return results

    async def _augment_all(
        trajs: list[dict[str, Any]],
        n: int,
        mdl: str,
    ) -> list[dict[str, Any]]:
        all_tasks = [_augment_one(t, provider, n, mdl) for t in trajs]
        nested = await asyncio.gather(*all_tasks)
        return [record for group in nested for record in group]

    augmented = asyncio.run(_augment_all(trajectories, n_variants, model))
    logger.info(
        "Augmented %d trajectories into %d records (n_variants=%d)",
        len(trajectories),
        len(augmented),
        n_variants,
    )
    return augmented

def split_by_task_id(
    records: list[dict[str, Any]],
    test_fraction: float = 0.2,
    seed: int = 42,
) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]:
    """Split records at task-ID boundary with no task_id crossing train/test.

    Augmented records that share a task_id are always assigned to the same
    partition, preventing task-family leakage.

    Args:
        records: List of record dicts, each with a 'task_id' field.
        test_fraction: Fraction of unique task_ids to assign to test set.
            Minimum 1 task_id goes to test even if fraction rounds to 0.
        seed: Random seed for reproducible splits.

    Returns:
        Tuple of (train_records, test_records) where task_ids never overlap.
    """
    task_ids = sorted({r["task_id"] for r in records})
    rng = random.Random(seed)
    rng.shuffle(task_ids)
    n_test = max(1, int(len(task_ids) * test_fraction))
    test_ids = set(task_ids[:n_test])
    train = [r for r in records if r["task_id"] not in test_ids]
    test = [r for r in records if r["task_id"] in test_ids]
    return train, test

def normalize_mined_pairs(
    trajectory: dict[str, Any],
    compress: bool = True,
    max_diff_lines: int = 500,
    language: str | None = None,
) -> list[dict[str, Any]]:
    """Convert a mined PR trajectory into per-step training pairs.

    Each review-to-revision cycle becomes one training record with
    activation_text (task + current code + review feedback) and
    teacher_text (activation + revision diff). Compatible with
    ``augment_trajectories``, ``split_by_task_id``, and ``save_jsonl``.

    The algorithm groups contiguous commits and reviews into blocks,
    then pairs each reviews-block with the following commits-block.
    Multiple commits in a block: the last commit is used (the state
    the reviewer actually saw). Multiple reviews: concatenated.

    Args:
        trajectory: Raw mined trajectory from ``mine_pr_diff_chains``.
        compress: Apply diff compression via ``compress_diff``.
        max_diff_lines: Max lines per compressed diff.
        language: Language tag for metadata (from repos config).

    Returns:
        List of training pair records with task_id, activation_text,
        teacher_text, and metadata fields.
    """
    task_id: str = trajectory.get("task_id", "")
    task_desc: str = trajectory.get("task_description", "")
    raw_steps: list[dict[str, Any]] = trajectory.get("steps", [])
    outcome: str = trajectory.get("outcome", "")

    if not raw_steps:
        return []

    blocks = _group_steps_into_blocks(raw_steps)

    def _diff(step: dict[str, Any]) -> str:
        raw = step.get("content", "")
        return compress_diff(raw, max_lines=max_diff_lines) if compress else raw

    def _record(idx: int, activation: str, teacher: str) -> dict[str, Any]:
        return _make_pair_record(task_id, outcome, language, idx, activation, teacher)

    records: list[dict[str, Any]] = []
    step_idx = 0
    prev_diff = ""
    bi = 0  # block index

    # --- Step 0: initial commits block ---
    if blocks[0][0] == "commit":
        last_commit = blocks[0][1][-1]
        diff = _diff(last_commit)
        activation = f"## Task\n{task_desc}"
        teacher = f"{activation}\n\n## Implementation\n{diff}"
        records.append(_record(step_idx, activation, teacher))
        prev_diff = diff
        step_idx += 1
        bi = 1

    # --- Subsequent (reviews, commits) pairs ---
    while bi < len(blocks) - 1:
        if blocks[bi][0] == "review" and blocks[bi + 1][0] == "commit":
            review_text = "\n\n".join(r.get("content", "") for r in blocks[bi][1])
            revision = _diff(blocks[bi + 1][1][-1])
            activation = f"## Task\n{task_desc}"
            if prev_diff:
                activation += f"\n\n## Current Code\n{prev_diff}"
            activation += f"\n\n## Review Feedback\n{review_text}"
            teacher = f"{activation}\n\n## Revision\n{revision}"
            records.append(_record(step_idx, activation, teacher))
            prev_diff = revision
            step_idx += 1
            bi += 2
        else:
            bi += 1

    return records

def compress_diff(content: str, max_lines: int = 500) -> str:
    """Filter irrelevant files and truncate large diffs.

    Parses the ``--- filename ---`` section format produced by
    ``mine_pr_diff_chains`` and removes lockfiles, generated code,
    binary files, and build artifacts.

    Args:
        content: Concatenated diff content with section headers.
        max_lines: Maximum total output lines.

    Returns:
        Filtered and truncated diff string.
    """
    if not content:
        return content

    parts = _SECTION_RE.split(content)

    # No section headers found -- fall through to truncation only
    if len(parts) == 1:
        return _truncate(content, max_lines)

    # parts = [preamble, header1, body1, header2, body2, ...]
    kept: list[str] = []
    idx = 1  # skip preamble (parts[0])
    while idx < len(parts):
        header = parts[idx]
        body = parts[idx + 1] if idx + 1 < len(parts) else ""
        filename = header[4:-4]  # strip "--- " and " ---"
        if not _should_skip(filename):
            kept.append(header + body)
        idx += 2

    result = "".join(kept).strip()

    if not result:
        return ""

    return _truncate(result, max_lines)

def apply_functional_lora(
    model: Any, lora_dict: dict[str, Any], hc: Any
) -> _FunctionalLoRAContext:
    """Create a context manager that patches model with functional LoRA.

    Usage:
        with apply_functional_lora(base_model, lora_dict, hc):
            output = base_model(input_ids)
            loss = criterion(output, target)
        loss.backward()  # gradients flow through A/B to hypernetwork

    Args:
        model: Base transformer model (nn.Module).
        lora_dict: Dict from HyperLoRA.generate_weights(). Structure:
            lora_dict[proj_name]["A"] shape: (batch=1, n_layers, r, d_in)
            lora_dict[proj_name]["B"] shape: (batch=1, n_layers, r, d_out)
            Keys match hc.lora_config.target_modules.
            Batch dimension is always 1 (squeezed at index 0).
        hc: HypernetConfig with attributes:
            hc.lora_config.target_modules: list of projection names
            hc.lora_config.r: LoRA rank
            hc.lora_config.lora_alpha: scaling numerator
            hc.layer_indices: list of absolute layer indices

    Returns:
        _FunctionalLoRAContext to use as context manager.
    """
    return _FunctionalLoRAContext(model, lora_dict, hc)