atropos/example_trainer/data.py

"""
Data processing utilities for GRPO trainer.

Handles data retrieval from Atropos API, padding, batching,
and advantage normalization.

Also extracts inference logprobs for alignment validation with training logprobs.
"""

import json
import math
import time
from typing import List, Optional, Tuple

import numpy as np
import torch

from .api import get_batch


def pad_data_to_good_offset(
    data: dict, 
    batch_size: int,
    extract_inference_logprobs: bool = True,
) -> Tuple[
    List[torch.Tensor], 
    List[torch.Tensor], 
    List[torch.Tensor], 
    List[torch.Tensor],
    Optional[List[np.ndarray]],
]:
    """
    Pad and batch data from the Atropos API.
    
    Processes raw batch data into properly padded tensors suitable for training:
    - Pads token sequences to nearest multiple of 64
    - Normalizes advantage scores
    - Extracts temperature values
    - Optionally extracts inference logprobs for alignment validation
    
    Args:
        data: Raw batch data from Atropos API
        batch_size: Size of each training batch
        extract_inference_logprobs: Whether to extract inference logprobs
        
    Returns:
        Tuple of (token_batches, label_batches, advantage_batches, temperature_batches, inference_logprobs)
        inference_logprobs is None if extract_inference_logprobs=False or no logprobs in data
    """
    max_token_len = max(
        [max([len(x) for x in item["tokens"]]) for item in data["batch"]]
    )
    
    # Pad to nearest multiple of 64 for GPU efficiency
    good_multiple = 64
    if (max_token_len - 1) % (good_multiple) != 0:
        max_token_len = math.ceil((max_token_len - 1) / (good_multiple)) * good_multiple
        token_setup_len = max_token_len + 1  # +1 for causal shift
    else:
        token_setup_len = max_token_len
        max_token_len = max_token_len - 1  # -1 for causal shift
    
    # Process all items
    input_ids = []
    labels = []
    advantages = []
    lengths = []
    temperatures = []
    inference_logprobs_list: List[np.ndarray] = []
    
    for item in data["batch"]:
        # Normalize advantage scores
        scores = np.array(item["scores"])
        if len(scores) > 1:
            scores = scores - scores.mean()
            scores = scores / max(scores.std(), 1e-8)
        item["scores"] = scores
        
        # Handle score overrides
        if item["overrides"] is not None:
            for i in range(len(item["overrides"])):
                if item["overrides"][i].get("set_advantage_to_zero", False):
                    item["scores"][i] = 0
        
        # Process each sample in the item
        for i in range(len(item["tokens"])):
            lengths.append(
                math.ceil((len(item["tokens"][i]) - 1) / good_multiple) * good_multiple
            )
            
            # Create labels with padding
            label_item = np.concatenate([
                np.array(item["masks"][i]),
                np.full(
                    max(0, token_setup_len - len(item["tokens"][i])),
                    -100,
                    dtype=np.int32,
                ),
            ])
            
            # Pad tokens
            item["tokens"][i] = np.concatenate([
                np.array(item["tokens"][i]),
                np.zeros(
                    max(0, token_setup_len - len(item["tokens"][i])),
                    dtype=np.int32,
                ),
            ])
            
            input_ids.append(item["tokens"][i][:-1])  # Remove last for causal
            labels.append(label_item[1:])  # Shift by 1 for causal
            advantages.append(item["scores"][i])
            
            # Extract inference logprobs for alignment validation
            # These come from vLLM during rollout generation
            if extract_inference_logprobs and "inference_logprobs" in item:
                if i < len(item["inference_logprobs"]):
                    inference_logprobs_list.append(
                        np.array(item["inference_logprobs"][i], dtype=np.float32)
                    )
            
            # Extract temperature (priority: override > generation_params > group_overrides > 1.0)
            t = 1.0
            if (
                item.get("overrides")
                and i < len(item["overrides"])
                and isinstance(item["overrides"][i], dict)
                and ("temperature" in item["overrides"][i])
            ):
                t = float(item["overrides"][i]["temperature"])
            elif item.get("generation_params") and ("temperature" in item["generation_params"]):
                t = float(item["generation_params"]["temperature"])
            elif item.get("group_overrides") and ("temperature" in item["group_overrides"]):
                t = float(item["group_overrides"]["temperature"])
            temperatures.append(t)
    
    # Batch the data
    token_batches = []
    label_batches = []
    advantage_batches = []
    temperature_batches = []
    
    for i in range(len(input_ids) // batch_size):
        start = i * batch_size
        end = (i + 1) * batch_size
        
        token_batches.append(
            torch.tensor(np.stack(input_ids[start:end], axis=0))
        )
        label_batches.append(
            torch.tensor(np.stack(labels[start:end], axis=0))
        )
        advantage_batches.append(
            torch.tensor(np.stack(advantages[start:end], axis=0)).view(-1, 1)
        )
        temperature_batches.append(
            torch.tensor(
                np.array(temperatures[start:end], dtype=np.float32)
            ).view(-1, 1, 1)
        )
    
    # Return inference logprobs if available
    inference_logprobs = inference_logprobs_list if inference_logprobs_list else None
    
    return token_batches, label_batches, advantage_batches, temperature_batches, inference_logprobs


def get_data(
    batch_size: int,
    seq_len: int,
    atropos_url: str = "http://localhost:8000",
    extract_inference_logprobs: bool = True,
) -> Tuple[
    List[Tuple[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[torch.Tensor]]],
    Optional[List[np.ndarray]],
]:
    """
    Fetch and process training data from the Atropos API.
    
    Continuously polls the API until data is available, then processes
    all available batches.
    
    Args:
        batch_size: Size of each training batch
        seq_len: Maximum sequence length (for reference, not used directly)
        atropos_url: URL of the Atropos API server
        extract_inference_logprobs: Whether to extract inference logprobs for alignment
        
    Returns:
        Tuple of (batches, all_inference_logprobs)
        - batches: List of processed batch tuples
        - all_inference_logprobs: List of inference logprob arrays for alignment validation
    """
    batches = []
    all_inference_logprobs: List[np.ndarray] = []
    
    while True:
        data = get_batch(url=atropos_url)
        
        if data["batch"] is not None:
            # Save batch for debugging
            with open("temp.json", "w", encoding="utf-8") as f:
                json.dump(data, f)
            
            # Process and accumulate batches
            token_batches, label_batches, adv_batches, temp_batches, inf_logprobs = \
                pad_data_to_good_offset(data, batch_size, extract_inference_logprobs)
            
            batches.append((token_batches, label_batches, adv_batches, temp_batches))
            
            if inf_logprobs:
                all_inference_logprobs.extend(inf_logprobs)
                
        elif len(batches) > 0:
            # Return accumulated batches when no more data
            return batches, all_inference_logprobs if all_inference_logprobs else None
        else:
            # Wait for data
            time.sleep(1)