adapters.py

from __future__ import annotations

import os
from typing import Any, Callable, Literal

import torch
from torch import Tensor
from torch.utils.data import Dataset
from transformers import PreTrainedTokenizerBase
from transformers import AutoModelForCausalLM
from transformers import PreTrainedModel

def tokenize_prompt_and_output(prompt_strs, output_strs, tokenizer):
    all_sequences = []
    mask_sequences = []
    
    for prompt, output in zip(prompt_strs, output_strs):
        prompt_tokens = tokenizer(prompt, add_special_tokens=False)['input_ids']
        output_tokens = tokenizer(output, add_special_tokens=False)['input_ids']
        
        full_sequence = prompt_tokens + output_tokens
        all_sequences.append(full_sequence)
        
        mask = [0] * len(prompt_tokens) + [1] * len(output_tokens)
        mask_sequences.append(mask)
    
    max_len = max(len(seq) for seq in all_sequences)
    
    padded_sequences = []
    padded_masks = []
    
    for seq, mask in zip(all_sequences, mask_sequences):
        padded_seq = seq + [tokenizer.pad_token_id] * (max_len - len(seq))
        padded_sequences.append(padded_seq)
        
        padded_mask = mask + [0] * (max_len - len(mask))
        padded_masks.append(padded_mask)
    
    input_ids = torch.tensor(padded_sequences, dtype=torch.long)
    response_mask = torch.tensor(padded_masks, dtype=torch.bool)
        
    return {
        'input_ids': input_ids[:, :-1],
        'labels': input_ids[:, 1:], 
        'response_mask': response_mask[:, 1:]
    }



def run_tokenize_prompt_and_output(
    prompt_strs: list[str],
    output_strs: list[str],
    tokenizer: PreTrainedTokenizerBase,
) -> dict[str, Tensor]:
    

    return (
        tokenize_prompt_and_output(prompt_strs, output_strs, tokenizer)
    )

def compute_group_normailzed_rewards(
    reward_fn: Callable,
    rollout_responses: list[str],
    repeated_ground_truths: list[str],
    group_size: int,
    advantage_eps: float,
    normalize_by_std: bool,
):

    raw_rewards = [reward_fn(prompt, gt)['reward'] for prompt, gt in zip(rollout_responses, repeated_ground_truths)]
    group_rewards = [raw_rewards[i:i + group_size] for i in range(0, len(raw_rewards), group_size)]

    advantages = []

    for i, group_reward in enumerate(group_rewards):
        group_reward = torch.tensor(group_reward)
        mean = torch.mean(group_reward)
        std = torch.std(group_reward)

        for reward in group_reward:
            advantage = reward - mean
            if normalize_by_std:
                advantage /= (std) + advantage_eps
            advantages.append(advantage)

    advantages = torch.Tensor(advantages)
    raw_rewards = torch.Tensor(raw_rewards)
    metadata = {}
    metadata['advantages'] = advantages
    metadata['raw_rewards'] = raw_rewards

    return (advantages, raw_rewards, metadata)





def run_compute_group_normalized_rewards(
    reward_fn: Callable,
    rollout_responses: list[str],
    repeated_ground_truths: list[str],
    group_size: int,
    advantage_eps: float,
    normalize_by_std: bool,
) -> tuple[torch.Tensor, dict[str, float]]:
    """
    Compute rewards for each group of rollout responses, 
    normalized by the group size.

    For more on GRPO, see:
        DeepSeekMath: https://arxiv.org/abs/2402.03300
        DeepSeek-R1: https://arxiv.org/abs/2501.12948

    Args:
        reward_fn: Callable[[str, str], dict[str, float]], 
            scores the rollout responses against the ground truths, 
            producing a dict with keys 
            "reward", "format_reward", and "answer_reward".
        rollout_responses: list[str], rollouts from the policy. 
            The length of this list is 
            `rollout_batch_size = n_prompts_per_rollout_batch * group_size`.
        repeated_ground_truths: list[str], the ground truths for the examples. 
            The length of this list is `rollout_batch_size`, 
            because the ground truth for each example is repeated `group_size` times.
        group_size: int, number of rollouts per group.
        advantage_eps: float, epsilon to avoid division by zero
            during group normalization.
        normalize_by_std: bool, whether to normalize the rewards by
            std(rewards).

    Returns:
        tuple[torch.Tensor, torch.Tensor, dict[str, float]]:
            torch.Tensor of shape (rollout_batch_size,): 
                group-normalized rewards for each rollout response.
            torch.Tensor of shape (rollout_batch_size,): 
                raw rewards for each rollout response.
            dict[str, float]: metadata for the rewards of the rollout batch.
                You may choose what you wish to log here
                (some statistics of the rewards, etc.).
    """
    return compute_group_normailzed_rewards(
    reward_fn,
    rollout_responses,
    repeated_ground_truths,
    group_size,
    advantage_eps,
    normalize_by_std)



def compute_entropy(logits: torch.Tensor) -> torch.Tensor:
    """
    Get the entropy of the logits (i.e., entropy of the final dimension).
    
    Args:
        logits: shape (..., vocab_size) - can be any shape with vocab_size as last dim
    
    Returns:
        entropy: shape (...) - same shape as input minus last dimension
    """
    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
    probs = torch.exp(log_probs)
    
    entropy = -torch.sum(probs * log_probs, dim=-1)
    return entropy


def run_compute_entropy(logits: torch.Tensor) -> torch.Tensor:
    """
    Adapter function that calls compute_entropy.
    This is what the test imports and calls.
    """
    return compute_entropy(logits)


def get_response_log_probs(
    model: PreTrainedModel,
    input_ids: torch.Tensor,
    labels: torch.Tensor,
    return_token_entropy: bool = False,
    ) -> dict[str, torch.Tensor]:

    logits = model(input_ids).logits

    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)
    log_probs = torch.gather(log_probs, dim=-1, index=labels.unsqueeze(-1))
    log_probs = log_probs.squeeze(-1)

    ret_dict = {}
    ret_dict['log_probs'] = log_probs

    if return_token_entropy:
        ret_dict['token_entropy'] = compute_entropy(logits)
    
    return ret_dict


def run_get_response_log_probs(
    model: torch.nn.Module,
    input_ids: torch.Tensor,
    labels: torch.Tensor,
    return_token_entropy: bool,
) -> dict[str, torch.Tensor]:
    """
    Adapter function that calls get_response_log_probs.
    This is what the test imports and calls.
    
    Args:
        model: PreTrainedModel HuggingFace model used for scoring
        input_ids: torch.Tensor shape (batch_size, sequence_length)
        labels: torch.Tensor shape (batch_size, sequence_length)
        return_token_entropy: bool - If True, also return per-token entropy
    
    Returns:
        dict[str, torch.Tensor] with "log_probs" and optionally "token_entropy"
    """
    return get_response_log_probs(
        model=model,
        input_ids=input_ids,
        labels=labels,
        return_token_entropy=return_token_entropy
    )

def compute_naive_policy_gradient_loss(
        raw_rewards_or_advantages,
        policy_log_probs
):
    pg_loss = -raw_rewards_or_advantages.expand_as(policy_log_probs) * policy_log_probs
    return pg_loss

def run_compute_naive_policy_gradient_loss(
    raw_rewards_or_advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
) -> torch.Tensor:
    """Compute policy gradient loss using either raw rewards or advantages.

    Args:
        raw_rewards_or_advantages: torch.Tensor of shape (batch_size, 1): 
            the raw rewards or advantages for each rollout response.
        policy_log_probs: torch.Tensor of shape (batch_size, sequence_length): 
            the log-probs of the policy.

    Returns:
        torch.Tensor of shape (batch_size, sequence_length): 
            the policy gradient per-token loss.
    """
    return compute_naive_policy_gradient_loss(raw_rewards_or_advantages, policy_log_probs)

def compute_grpo_clip_loss(
        advantages,
        policy_log_probs,
        old_log_probs,
        cliprange
):
    advantages = advantages.expand_as(policy_log_probs)

    improvement = torch.exp(policy_log_probs - old_log_probs)

    arg_1 = improvement * advantages

    arg_2 = torch.clip(improvement, 1.0-cliprange, 1.0+cliprange) * advantages

    loss = -torch.min(arg_1, arg_2)


    clip_mask = (arg_1 != arg_2)

    clipped_percentage = clip_mask.sum() / arg_2.numel()

    metadata = {}

    metadata["clipped_percentage"] = clipped_percentage


    return (loss, metadata)


def run_compute_grpo_clip_loss(
    advantages: torch.Tensor,
    policy_log_probs: torch.Tensor,
    old_log_probs: torch.Tensor,
    cliprange: float,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """Compute the GRPO-Clip loss.

    Args:
        advantages: torch.Tensor of shape (batch_size, 1): 
            the advantages for each rollout response.
        policy_log_probs: torch.Tensor of shape (batch_size, sequence_length): 
            the log-probs of the policy.
        old_log_probs: torch.Tensor of shape (batch_size, sequence_length): 
            the log-probs of the old policy.
        cliprange: float, the clip range for the ratio.

    Returns:
        tuple[torch.Tensor, dict[str, torch.Tensor]]:
            torch.Tensor of shape (batch_size, sequence_length): 
                the GRPO-Clip per-token loss.
            dict[str, torch.Tensor]: metadata for the GRPO-Clip loss 
                (used to compute clip fraction).
    """
    return compute_grpo_clip_loss(advantages, policy_log_probs, old_log_probs, cliprange)

def compute_policy_gradient_loss(
    policy_log_probs,
    loss_type,
    raw_rewards,
    advantages,
    old_log_probs,
    cliprange
):
    if loss_type == 'no_baseline':
        loss = compute_naive_policy_gradient_loss(raw_rewards, policy_log_probs)
        metadata = {}
    
    elif loss_type == 'reinforce_with_baseline':
        loss = compute_naive_policy_gradient_loss(advantages, policy_log_probs)
        metadata = {}
    
    elif loss_type == 'grpo_clip':
        loss, metadata = compute_grpo_clip_loss(advantages, policy_log_probs, old_log_probs, cliprange)
    
    return (loss, metadata)
    



def run_compute_policy_gradient_loss(
    policy_log_probs: torch.Tensor,
    loss_type: str,
    raw_rewards: torch.Tensor,
    advantages: torch.Tensor,
    old_log_probs: torch.Tensor,
    cliprange: float,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """
    Wrapper that delegates to the appropriate policy gradient loss function above.
    """
    return compute_policy_gradient_loss(policy_log_probs, loss_type, raw_rewards, advantages, old_log_probs, 
                                        cliprange)

def masked_mean(tensor, mask, dim):
    if dim is None:
        mean = torch.mean(tensor[mask == 1])
    else:
        masked = torch.where(mask == 1, tensor, 0.0)
        mean = torch.sum(masked, dim=dim) /torch.sum(mask == 1, dim=dim)
    return mean


def run_masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int | None = None) -> torch.Tensor:
    """Compute the mean of the tensor along a dimension,
    considering only the elements with mask value 1.

    Args:
        tensor: torch.Tensor, the tensor to compute the mean of.
        mask: torch.Tensor, the mask. We only take the mean over
            the elements with mask value 1.
        dim: int | None, the dimension to compute the mean along.
            If None, sum over all non-masked elements and average
            by their total count.

    Returns:
        torch.Tensor, the mean of the tensor along the specified
            dimension, considering only the elements with mask value 1.
    """
    return masked_mean(tensor, mask, dim)




def sft_microbatch_train_step(
    policy_log_probs: torch.Tensor,
    response_mask: torch.Tensor,
    gradient_accumulation_steps: int,
    normalize_constant: float = 1.0,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """
    Execute a forward-and-backward pass on a microbatch.
    
    Args:
        policy_log_probs: (batch_size, sequence_length), per-token log-probabilities from the 
                         SFT policy being trained.
        response_mask: (batch_size, sequence_length), 1 for response tokens, 0 for prompt/padding.
        gradient_accumulation_steps: Number of microbatches per optimizer step.
        normalize_constant: The constant by which to divide the sum. It is fine to leave this as 1.0.
    
    Returns:
        tuple[torch.Tensor, dict[str, torch.Tensor]]:
            - loss: scalar tensor. The microbatch loss, adjusted for gradient accumulation.
            - metadata: Dict with metadata from the underlying loss call, and any other statistics.
    """

    response_nll = -run_masked_normalize(policy_log_probs, response_mask, normalize_constant, dim=-1).mean()

    microbatch_loss = response_nll / gradient_accumulation_steps

    microbatch_loss.backward()

    metadata = {}
    metadata['loss'] = microbatch_loss.item()  

    return (microbatch_loss, metadata)


def run_sft_microbatch_train_step(
    policy_log_probs: torch.Tensor,
    response_mask: torch.Tensor,
    gradient_accumulation_steps: int,
    normalize_constant: float = 1.0,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """
    Adapter function that calls sft_microbatch_train_step.
    This is what the test imports and calls.
    """
    return sft_microbatch_train_step(
        policy_log_probs=policy_log_probs,
        response_mask=response_mask,
        gradient_accumulation_steps=gradient_accumulation_steps,
        normalize_constant=normalize_constant
    )

def grpo_microbatch_train_step(
        policy_log_probs,
        response_mask,
        gradient_accumulation_steps,
        loss_type,
        raw_rewards,
        advantages,
        old_log_probs,
        clip_range
):
    loss, metadata = compute_policy_gradient_loss(policy_log_probs, loss_type, raw_rewards, advantages, old_log_probs,
                                           clip_range)
    
    loss = masked_mean(loss, response_mask, dim=None)

    loss /= gradient_accumulation_steps

    loss.backward()


    return loss, metadata

    
def run_grpo_microbatch_train_step(
    policy_log_probs: torch.Tensor,
    response_mask: torch.Tensor,
    gradient_accumulation_steps: int,
    loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"],
    raw_rewards: torch.Tensor | None = None,
    advantages: torch.Tensor | None = None,
    old_log_probs: torch.Tensor | None = None,
    cliprange: float | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """Compute the policy gradient loss and backprop its gradients for a microbatch.

    Args:
        policy_log_probs: torch.Tensor of shape (batch_size, sequence_length): 
            the log-probs of the policy.
        response_mask: torch.Tensor of shape (batch_size, sequence_length): 
            the mask for the response.
        gradient_accumulation_steps: int, the number of gradient accumulation steps.
        loss_type: Literal["no_baseline", "reinforce_with_baseline", "grpo_clip"], 
            the type of loss function to use.
        raw_rewards: torch.Tensor | None, the raw rewards for each rollout response.
            Needed for loss_type="no_baseline".
        advantages: torch.Tensor | None, the advantages for each rollout response.
            Needed for loss_type in {"reinforce_with_baseline", "grpo_clip"}.
        old_log_probs: torch.Tensor | None, the log-probs of the old policy.
            Needed for loss_type="grpo_clip".
        cliprange: float | None, the clip range for the ratio. 
            Needed for loss_type="grpo_clip".
        constant_normalize_factor: int | None, provided if we want to sum over 
            the sequence dimension and normalize by this constant factor
            (as in Dr. GRPO).

    Returns:
        tuple[torch.Tensor, dict[str, torch.Tensor]]: 
            the policy gradient loss and its metadata.
    """
    return grpo_microbatch_train_step(
        policy_log_probs,
        response_mask,
        gradient_accumulation_steps,
        loss_type,
        raw_rewards,
        advantages,
        old_log_probs,
        cliprange
)



def run_masked_normalize(
    tensor: torch.Tensor,
    mask: torch.Tensor,
    normalize_constant: float,
    dim: int | None = None,
) -> torch.Tensor:
    """Sum over a dimension and normalize by a constant,
    considering only the elements with mask value 1.
    """
    masked_tensor = tensor * mask.float()
    
    if dim is None:
        masked_sum = torch.sum(masked_tensor)
    else:
        masked_sum = torch.sum(masked_tensor, dim=dim)
    
    return masked_sum / normalize_constant


"""
The below adapters are used in the optional 
RLHF / safety part of the Alignment assignment.
"""


def get_packed_sft_dataset(
    tokenizer: PreTrainedTokenizerBase,
    dataset_path: str | os.PathLike,
    seq_length: int,
    shuffle: bool,
) -> Dataset:
    """
    Given a tokenizer and a path to a dataset with instruction-tuning examples,
    construct a PyTorch Dataset for language modeling. The examples should be
    packed, i.e., all sequences in the dataset are of a constant length (`seq_length`).

    Args:
        tokenizer: transformers.PreTrainedTokenizerBase
            Transformers tokenizer to use in tokenizing and encoding text.
        dataset_path: str
            Path to file with instruction-tuning examples.
        seq_length: int
            Number of tokens to include in each example.
        shuffle: bool
            If true, shuffle the documents before packing them into examples.

    Returns:
        PyTorch Dataset for language modeling. Each example in this dataset is a dictionary of
        with keys "input_ids" and "labels" (both tensors of shape (seq_length, )).
        "input_ids" contains the token IDs for the language modeling inputs, and "labels" contains
        the token IDs for the language modeling labels.
    """
    raise NotImplementedError


def run_iterate_batches(
    dataset: Dataset,
    batch_size: int,
    shuffle: bool,
):
    """
    Given a PyTorch Dataset, return an iterable over batches of size `batch_size`.
    Iterating through the returned iterable should constitute one epoch over the Dataset.

    Args:
        dataset: Dataset
            Dataset to emit batches from.
        batch_size: int
            Number of examples to include per batch.
        shuffle: bool
            If true, shuffle examples before batching them.

    Returns:
        Iterable over batches, where each batch has size `batch_size`.
    """
    raise NotImplementedError


def run_parse_mmlu_response(
    mmlu_example: dict[str, Any],
    model_output: str,
) -> str | None:
    """
    Given an MMLU example and a model output, parse the model output into a
    predicted option letter (i.e., 'A', 'B', 'C', or 'D'). If the model output
    cannot be parsed into a prediction option letter, return None.

    mmlu_example: dict[str, Any]
        Dictionary with an MMLU example. Contains the following keys:
        - "subject": str with the subject of the question.
        - "question": str with the text of the question.
        - "options": list[str] with the four answer options (in order).
                     The first option refers to letter "A", the second to "B", etc.
        - "answer": str with the option of the correct answer (e.g., "A")
    model_output: str
        str with the model's output to the MMLU example.

    Returns:
        str (one of "A", "B", "C", or "D") if the model output can be parsed into a prediction,
        else None.
    """
    raise NotImplementedError


def run_parse_gsm8k_response(
    model_output: str,
) -> str | None:
    """
    Given a GSM8K model output, parse the model output into a predicted numeric answer by
    taking the last number that occurs in the output.

    model_output: str
        str with the model's output to a GSM8K example.

    Returns:
        str with the predicted numeric answer if the model output can be parsed into a prediction,
        else None.
    """
    raise NotImplementedError


def run_compute_per_instance_dpo_loss(
    lm: torch.nn.Module,
    lm_ref: torch.nn.Module,
    tokenizer: PreTrainedTokenizerBase,
    beta: float,
    prompt: str,
    response_chosen: str,
    response_rejected: str,
) -> torch.Tensor:
    """
    Given two language models (`lm`, and the "reference model" `lm_ref`),
    their tokenizer, the DPO beta hyperparameter, a prompt and a pair
    of responses to the prompt, computes the value of the DPO loss for this example.

    lm: torch.nn.Module
        Language model being trained.
    lm_ref: torch.nn.Module
        Reference language model.
    tokenizer: PreTrainedTokenizerBase
        Tokenizer for both language models.
    beta: float
        DPO beta hyperparameter.
    prompt: str
        Prompt for this instance of preference pair.
    response_chosen: str
        Preferred response to the prompt.
    response_rejected: str
        Rejected response to the prompt.

    Returns:
        torch.Tensor with the DPO loss for this example.
    """
    raise NotImplementedError