models/FarSLIP/open_clip_train/train.py

import json
import logging
import math
import os
import time

import numpy as np
import torch
import torch.nn.functional as F
from torch.nn.parallel.distributed import DistributedDataParallel

try:
    import wandb
except ImportError:
    wandb = None

from open_clip import get_input_dtype, CLIP, CustomTextCLIP
from open_clip_train.distributed import is_master
from open_clip_train.zero_shot import zero_shot_eval
from open_clip_train.precision import get_autocast

from torchvision.ops import roi_align
from torchvision.transforms import RandomResizedCrop

class AverageMeter(object):
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()
        self.total_sum =0

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count
        self.total_sum += val * n


def postprocess_clip_output(model_out):
    return {
        "image_features": model_out[0],
        "text_features": model_out[1],
        "logit_scale": model_out[2]
    }


def unwrap_model(model):
    if hasattr(model, 'module'):
        return model.module
    else:
        return model


def backward(total_loss, scaler):
    if scaler is not None:
        scaler.scale(total_loss).backward()
    else:
        total_loss.backward()


def filter_batch_data(batch, used_losses, input_dtype, device):
    selected = {}

    selected["global_image"] = batch["global_image"].to(device=device, dtype=input_dtype, non_blocking=True)
    if 'global_itc' in used_losses:
        selected["global_text"] = batch["global_text"].to(device=device, non_blocking=True)
    if 'local_itc' in used_losses:
        selected["local_images"] = batch["local_images"].to(device=device, dtype=input_dtype, non_blocking=True)
        selected["local_texts"] = batch["local_texts"].to(device=device, non_blocking=True)
        selected["boxes"] = batch["boxes"].to(device=device, non_blocking=True)
        if 'local_categories' in batch:
            selected["local_categories"] = batch["local_categories"].to(device=device)
    if 'distill' in used_losses and 'subset_images' in batch:
        selected["subset_images"] = batch["subset_images"].to(device=device, dtype=input_dtype, non_blocking=True)
        selected["subset_boxes"] = batch["subset_boxes"].to(device=device, non_blocking=True)

    return selected


def get_batch_data(batch, used_losses, input_dtype, device):
    images, global_caption = batch

    global_image = torch.stack([item["global"] for item in images], dim=0)  # (B, C, H, W)
    local_imgs = torch.stack([item["locals"] for item in images], dim=0)  # (B, N, C, H_l, W_l)
    bboxes = torch.stack([item["bboxes"] for item in images], dim=0)  # (B, N, 5)

    selected = {}
    selected["global_image"] = global_image.to(device=device, dtype=input_dtype, non_blocking=True)
    selected["global_text"] = global_caption.to(device=device, non_blocking=True)
    if 'distill' in used_losses:
        selected["subset_images"] = local_imgs.to(device=device, dtype=input_dtype, non_blocking=True)
        selected["subset_boxes"] = bboxes.to(device=device, non_blocking=True)

    return selected


def train_one_epoch(model, teacher, method, data, loss, mpcl_loss, epoch, optimizer, scaler, scheduler, args, tb_writer=None):
    device = torch.device(args.device)
    autocast = get_autocast(args.precision, device_type=device.type)
    input_dtype = get_input_dtype(args.precision)

    model.train()
    if teacher: teacher.eval()

    data['train'].set_epoch(epoch)  # set epoch in process safe manner via sampler or shared_epoch
    dataloader = data['train'].dataloader
    num_batches_per_epoch = dataloader.num_batches // args.accum_freq
    sample_digits = math.ceil(math.log(dataloader.num_samples + 1, 10))

    if args.accum_freq > 1:
        accum_images, accum_texts, accum_features = [], [], {}

    losses_m = {}
    batch_time_m = AverageMeter()
    data_time_m = AverageMeter()
    end = time.time()
    for i, batch in enumerate(dataloader):
        i_accum = i // args.accum_freq      # Accumulated gradient updates at batch (i)
        step = num_batches_per_epoch * epoch + i_accum      # Total steps for whole training

        if not args.skip_scheduler:
            scheduler(step)

        data_time_m.update(time.time() - end)
        optimizer.zero_grad()

        if method.startswith('farslip'):
            assert args.accum_freq == 1, "accum freq disabled"

            # ========== Obtain input batches ==========
            if method == 'farslip1':
                batch = get_batch_data(batch, input_dtype=input_dtype, used_losses=args.loss_type, device=device)
            elif method == 'farslip2':
                batch = filter_batch_data(batch, input_dtype=input_dtype, used_losses=args.loss_type, device=device)
            else: raise ValueError(f"Unknown args.method: {args.method}")

            # ========== Obtain VALID local inputs (images, boxes, categories) ==========
            if "local_itc" in args.loss_type:

                bboxes = batch["boxes"]  # (B, num_boxes, 5)
                local_images = batch["local_images"]  # (B, num_boxes, C, H, W)
                local_texts = batch["local_texts"]  # (B, num_boxes, context_length)

                local_texts_list = []
                bboxes_list, local_imgs_list = [], []
                local_categories_list = []  # (B, num_boxes)

                has_local_categories = 'local_categories' in batch
                if has_local_categories: local_categories = batch['local_categories']

                for idx in range(len(bboxes)):
                    bbox = bboxes[idx]
                    local_img = local_images[idx]
                    validity = bbox[:, -1] == 1
                    bboxes_list.append(bbox[validity, :4])
                    local_imgs_list.append(local_img[validity])
                    local_text = local_texts[idx]
                    local_texts_list.append(local_text[validity])

                    if has_local_categories: local_categories_list.append(local_categories[idx][validity])

                local_texts = torch.cat(local_texts_list)  # (valid_objects, context_length)
                batch["local_texts"] = local_texts

                local_images = torch.cat(local_imgs_list)  # (valid_objects, C, H, W)
                batch['local_images'] = local_images

                if has_local_categories: local_categories = torch.cat(local_categories_list)

            if "distill" in args.loss_type and "subset_images" in batch:

                batch['subset_images'] = batch['subset_images'].reshape(-1, *batch['subset_images'].shape[2:])
                subset_bboxes_list = [bbox[:, :4] for bbox in batch["subset_boxes"]]

            # ============= Start training ==============

            def pad_tensor(tensor, total_len):
                valid_len, feat_dim = tensor.shape
                padded = torch.zeros((total_len, feat_dim), device=tensor.device, dtype=tensor.dtype)
                padded[:valid_len] = tensor
                return padded, valid_len, feat_dim

            visual_backbone = model.module.visual if hasattr(model, 'module') else model.visual
            grid_size = visual_backbone.grid_size   

            with autocast():
                features = model(batch=batch, device=device, input_dtype=input_dtype, used_losses=args.loss_type, last_attn_type=args.last_attn_type)
                if teacher and "distill" in args.loss_type:
                    with torch.no_grad():
                        features_t = teacher(batch=batch, device=device, input_dtype=input_dtype, used_losses=args.loss_type, last_attn_type=args.last_attn_type)
                logit_scale = features["logit_scale"]

                def _denormalize_boxes(normed_boxes, x):
                    h, w = x.shape[-2:]
                    denormed_boxes = []
                    for boxes in normed_boxes:
                        new_boxes = boxes.clone()
                        new_boxes[:, [0, 2]] *= w
                        new_boxes[:, [1, 3]] *= h
                        denormed_boxes.append(new_boxes)
                    return denormed_boxes

                def extract_normed_roi_features(global_patches, bboxes_list, grid_size):
                    """
                    Extract ROI features from global_patches, then perform mean pooling and normalize.
                    """
                    B, N, D = global_patches.shape  # (B, N_patches, D)
                    H, W = grid_size

                    patches_2d = global_patches.view(B, H, W, D).permute(0, 3, 1, 2)  # (B, D, H, W)
                    rois = _denormalize_boxes(bboxes_list, patches_2d)  # (B*n, 5)
                    roi_feats = roi_align(patches_2d, rois, output_size=grid_size, spatial_scale=1.0, sampling_ratio=-1,
                                          aligned=True)
                    # roi_global_features = F.normalize(roi_feats, dim=-1)
                    pooled = roi_feats.mean(dim=[2, 3])                         # (B*n, D)
                    normed = F.normalize(pooled, dim=-1)
                    return normed

                # ============== feat from global patches ===============
                if "distill" in args.loss_type:
                    normed_pooled_roi_global_features = extract_normed_roi_features(
                        features["global_patches"], subset_bboxes_list, grid_size)

                losses = {}
                # ============ distill loss ============
                if "distill" in args.loss_type:     # ["roi2cls", "roi2pooled", "combined"]
                    if args.distill_align == 'roi2pooled':      # patch-to-patch alignment
                        if args.distill_type == 'active':
                            local_patches_t = features["subset_patches"]
                        else: local_patches_t = features_t["subset_patches"]

                        local_patches_t = local_patches_t.view(local_patches_t.shape[0], grid_size[0], grid_size[1],
                                                           -1).permute(0, 3, 1, 2)
                        pooled_local_patches_t = local_patches_t.mean(dim=[2, 3])
                        normed_pooled_local_patches_t = F.normalize(pooled_local_patches_t, dim=-1)
                        loss_cosine = 1.0 - (normed_pooled_roi_global_features * normed_pooled_local_patches_t).sum(-1).mean()    # global roi & local roi

                    elif args.distill_align == 'roi2cls':       # indirect alignment with text
                        if args.distill_type=='active':
                            normed_local_pooled_t = F.normalize(features["subset_image_pooled"], dim=-1)
                        else:
                            normed_local_pooled_t = F.normalize(features_t["subset_image_pooled"], dim=-1)
                        loss_cosine = 1.0 - (normed_pooled_roi_global_features * normed_local_pooled_t).sum(-1).mean()  # global roi & local pooled

                    elif args.distill_align == 'combined':
                        if args.distill_type == 'active':
                            local_patches_t = features["subset_patches"]
                            local_image_pooled_t = features["subset_image_pooled"]
                        else:
                            local_patches_t = features_t["subset_patches"]
                            local_image_pooled_t = features_t["subset_image_pooled"]

                        local_patches_t = local_patches_t.view(local_patches_t.shape[0], grid_size[0], grid_size[1],
                                                           -1).permute(0, 3, 1, 2)
                        pooled_local_patches_t = local_patches_t.mean(dim=[2, 3])
                        normed_pooled_local_patches_t = F.normalize(pooled_local_patches_t, dim=-1)

                        normed_local_pooled_t = F.normalize(local_image_pooled_t, dim=-1)

                        loss_roi2roi = 1.0 - (normed_pooled_roi_global_features * normed_pooled_local_patches_t).sum(-1).mean()
                        loss_roi2pooled = 1.0 - (normed_pooled_roi_global_features * normed_local_pooled_t).sum(-1).mean()
                        loss_cosine = 0.5 * loss_roi2roi + 0.5 * loss_roi2pooled
                    else:
                        raise ValueError(f"Unknown distill_align: {args.distill_align}")
                    losses["distill"] = loss_cosine

                # ============ local_itc loss ============
                if "local_itc" in args.loss_type:       # ["cls", "pooled", "roi"]
                    total_len = args.batch_size * args.max_boxes

                    if args.local_itc_align == "cls":
                        normed_local_pooled = F.normalize(features["local_image_pooled"], dim=-1)
                        local_feat = normed_local_pooled

                    elif args.local_itc_align == "pooled":
                        local_patches = features["local_patches"]
                        local_patches = local_patches.view(local_patches.shape[0], grid_size[0], grid_size[1],
                                                           -1).permute(0, 3, 1, 2)
                        pooled_local_patches = local_patches.mean(dim=[2, 3])
                        normed_pooled_local_patches = F.normalize(pooled_local_patches, dim=-1)

                        local_feat = normed_pooled_local_patches

                    elif args.local_itc_align == "roi":
                        local_feat = extract_normed_roi_features(
                            features["global_patches"], bboxes_list, grid_size)

                    else:
                        raise ValueError(f"Unknown local_itc_align: {args.local_itc_align}")

                    padded_normed_local_pooled, valid_len, feature_len = pad_tensor(local_feat, total_len)

                    normed_local_text_pooled = F.normalize(features["local_text_pooled"], dim=-1)
                    padded_normed_local_text_pooled, _, _ = pad_tensor(normed_local_text_pooled, total_len)

                    padding_mask = torch.zeros(total_len, dtype=torch.bool, device=device)
                    padding_mask[:valid_len] = True

                    padded_local_categories = torch.full((total_len,), -1, dtype=local_categories.dtype,
                                                         device=local_categories.device)
                    padded_local_categories[:valid_len] = local_categories

                    if mpcl_loss:
                        region_itc_loss = mpcl_loss(padded_normed_local_pooled, padded_normed_local_text_pooled, logit_scale, padded_local_categories)
                    else:
                        region_itc_loss = loss(padded_normed_local_pooled, padded_normed_local_text_pooled, logit_scale, padding_mask=padding_mask)

                    losses["region_itc"] = region_itc_loss

                # ============ global_itc loss ============
                if "global_itc" in args.loss_type:
                    normed_global_pooled = F.normalize(features["global_image_pooled"], dim=-1)
                    normed_global_text_pooled = F.normalize(features["global_text_pooled"], dim=-1)
                    assert normed_global_pooled.shape[0] == normed_global_text_pooled.shape[0]
                    global_itc_loss = loss(normed_global_pooled, normed_global_text_pooled, logit_scale)
                    losses["global_itc"] = global_itc_loss

                total_loss = losses.get("distill", 0) * args.w_d + losses.get("region_itc", 0) * args.w_l + losses.get("global_itc", 0) * args.w_g
                losses["loss"] = total_loss

            backward(total_loss, scaler)


        elif method == 'vanilla':
            # images, texts = batch
            images, texts = batch["global_image"], batch["global_text"]
            images = images.to(device=device, dtype=input_dtype, non_blocking=True)
            texts = texts.to(device=device, non_blocking=True)
            if args.accum_freq == 1:
                with autocast():
                    model_out = model(images, texts)
                    logit_scale = model_out["logit_scale"]
                    losses = loss(**model_out, output_dict=True)

                    total_loss = sum(losses.values())
                    losses["loss"] = total_loss

                backward(total_loss, scaler)
            else:
                # First, cache the features without any gradient tracking.
                with torch.no_grad():
                    with autocast():
                        model_out = model(images, texts)

                        for f in ("logit_scale", "logit_bias"):
                            model_out.pop(f, None)

                        for key, val in model_out.items():
                            if key in accum_features:
                                accum_features[key].append(val)
                            else:
                                accum_features[key] = [val]

                    accum_images.append(images)
                    accum_texts.append(texts)

                # If (i + 1) % accum_freq is not zero, move on to the next batch.
                if ((i + 1) % args.accum_freq) > 0:
                    # FIXME this makes data time logging unreliable when accumulating
                    continue

                # Now, ready to take gradients for the last accum_freq batches.
                # Re-do the forward pass for those batches, and use the cached features from the other batches as negatives.
                # Call backwards each time, but only step optimizer at the end.
                optimizer.zero_grad()
                for j in range(args.accum_freq):
                    images = accum_images[j]
                    texts = accum_texts[j]
                    with autocast():
                        model_out = model(images, texts)

                        inputs_no_accum = {}
                        inputs_no_accum["logit_scale"] = logit_scale = model_out.pop("logit_scale")
                        if "logit_bias" in model_out:
                            inputs_no_accum["logit_bias"] = model_out.pop("logit_bias")

                        inputs = {}
                        for key, val in accum_features.items():
                            accumulated = accum_features[key]
                            inputs[key] = torch.cat(accumulated[:j] + [model_out[key]] + accumulated[j + 1:])

                        losses = loss(**inputs, **inputs_no_accum, output_dict=True)
                        del inputs
                        del inputs_no_accum
                        total_loss = sum(losses.values())
                        losses["loss"] = total_loss

                    backward(total_loss, scaler)

        if teacher and args.distill_type == 'ema':
            # EMA update for the teacher
            momentum = args.EMA_momentum
            with torch.no_grad():
                for param_q, param_k in zip(model.parameters(), teacher.parameters()):
                    param_k.data.mul_(momentum).add_(
                        (1 - momentum) * param_q.detach().data)

        if scaler is not None:
            if args.horovod:
                optimizer.synchronize()
                scaler.unscale_(optimizer)
                if args.grad_clip_norm is not None:
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_norm, norm_type=2.0)
                with optimizer.skip_synchronize():
                    scaler.step(optimizer)
            else:
                if args.grad_clip_norm is not None:
                    scaler.unscale_(optimizer)
                    torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_norm, norm_type=2.0)
                scaler.step(optimizer)
            scaler.update()
        else:
            if args.grad_clip_norm is not None:
                torch.nn.utils.clip_grad_norm_(model.parameters(), args.grad_clip_norm, norm_type=2.0)
            optimizer.step()

        # reset gradient accum, if enabled
        if args.accum_freq > 1:
            accum_images, accum_texts, accum_features = [], [], {}

        # Note: we clamp to 4.6052 = ln(100), as in the original paper.
        with torch.no_grad():
            unwrap_model(model).logit_scale.clamp_(0, math.log(100))
            if teacher and args.distill_type == 'ema': unwrap_model(teacher).logit_scale.clamp_(0, math.log(100))

        batch_time_m.update(time.time() - end)
        end = time.time()
        batch_count = i_accum + 1
        if is_master(args) and (i_accum % args.log_every_n_steps == 0 or batch_count == num_batches_per_epoch):
            batch_size = dataloader.batch_size
            num_samples = batch_count * batch_size * args.accum_freq * args.world_size
            samples_per_epoch = dataloader.num_samples
            percent_complete = 100.0 * batch_count / num_batches_per_epoch

            # NOTE loss is coarsely sampled, just master node and per log update
            for key, val in losses.items():
                if key not in losses_m:
                    losses_m[key] = AverageMeter()
                losses_m[key].update(val.item(), batch_size)

            logit_scale_scalar = logit_scale.item()
            loss_log = " ".join(
                [
                    f"{loss_name.capitalize()}: {loss_m.val:#.5g} ({loss_m.avg:#.5g})"
                    for loss_name, loss_m in losses_m.items()
                ]
            )
            samples_per_second = args.accum_freq * args.batch_size * args.world_size / batch_time_m.val
            samples_per_second_per_gpu = args.accum_freq * args.batch_size / batch_time_m.val
            eta_total = batch_time_m.total_sum / (i+1) * (args.epochs - epoch - (i+1)/dataloader.num_batches) * dataloader.num_batches
            logging.info(
                f"Train Epoch: {epoch} [{num_samples:>{sample_digits}}/{samples_per_epoch} ({percent_complete:.0f}%)] "
                # f"ETA: {eta_total:.3f}s "
                f"ETA: {int(eta_total // 3600):02}h{int((eta_total % 3600) // 60):02}m{int(eta_total % 60):02}s "
                f"Data (t): {data_time_m.avg:.3f} "
                f"Batch (t): {batch_time_m.avg:.3f}, {samples_per_second:#g}/s, {samples_per_second_per_gpu:#g}/s/gpu "
                f"LR: {optimizer.param_groups[0]['lr']:5f} "
                f"Logit Scale: {logit_scale_scalar:.3f} " + loss_log
            )

            # Save train loss / etc. Using non avg meter values as loggers have their own smoothing
            log_data = {
                "data_time": data_time_m.val,
                "batch_time": batch_time_m.val,
                "samples_per_second": samples_per_second,
                "samples_per_second_per_gpu": samples_per_second_per_gpu,
                "scale": logit_scale_scalar,
                "lr": optimizer.param_groups[0]["lr"]
            }
            log_data.update({name:val.val for name,val in losses_m.items()})

            log_data = {"train/" + name: val for name, val in log_data.items()}

            if tb_writer is not None:
                for name, val in log_data.items():
                    tb_writer.add_scalar(name, val, step)

            if args.wandb:
                assert wandb is not None, 'Please install wandb.'
                log_data['step'] = step  # for backwards compatibility
                wandb.log(log_data, step=step)

            # resetting batch / data time meters per log window
            batch_time_m.reset()
            data_time_m.reset()

    # end for


def evaluate(model, data, epoch, args, tb_writer=None, tokenizer=None):
    metrics = {}
    if not is_master(args):
        return metrics
    device = torch.device(args.device)
    model.eval()

    zero_shot_metrics = zero_shot_eval(model, data, epoch, args, tokenizer=tokenizer)
    metrics.update(zero_shot_metrics)

    autocast = get_autocast(args.precision, device_type=device.type)
    input_dtype = get_input_dtype(args.precision)

    if 'val' in data and (args.val_frequency and ((epoch % args.val_frequency) == 0 or epoch == args.epochs)):
        dataloader = data['val'].dataloader
        num_samples = 0
        samples_per_val = dataloader.num_samples

        # FIXME this does not scale past small eval datasets
        # all_image_features @ all_text_features will blow up memory and compute very quickly
        cumulative_loss = 0.0
        cumulative_gen_loss = 0.0
        all_image_features, all_text_features = [], []
        with torch.inference_mode():
            for i, batch in enumerate(dataloader):
                images, texts = batch
                images = images.to(device=device, dtype=input_dtype, non_blocking=True)
                texts = texts.to(device=device, non_blocking=True)

                with autocast():
                    model_out = model(images, texts)
                    image_features = model_out["image_features"]
                    text_features = model_out["text_features"]
                    logit_scale = model_out["logit_scale"]
                    # features are accumulated in CPU tensors, otherwise GPU memory exhausted quickly
                    # however, system RAM is easily exceeded and compute time becomes problematic
                    all_image_features.append(image_features.cpu())
                    all_text_features.append(text_features.cpu())
                    logit_scale = logit_scale.mean()
                    logits_per_image = logit_scale * image_features @ text_features.t()
                    logits_per_text = logits_per_image.t()

                    batch_size = images.shape[0]
                    labels = torch.arange(batch_size, device=device).long()
                    total_loss = (
                        F.cross_entropy(logits_per_image, labels) +
                        F.cross_entropy(logits_per_text, labels)
                    ) / 2

                    gen_loss = maybe_compute_generative_loss(model_out)

                cumulative_loss += total_loss * batch_size
                num_samples += batch_size
                if is_master(args) and (i % 100) == 0:
                    logging.info(
                        f"Eval Epoch: {epoch} [{num_samples} / {samples_per_val}]\t"
                        f"Clip Loss: {cumulative_loss / num_samples:.6f}\t")

                    if gen_loss is not None:
                        cumulative_gen_loss += gen_loss * batch_size
                        logging.info(
                            f"Generative Loss: {cumulative_gen_loss / num_samples:.6f}\t")

            val_metrics = get_clip_metrics(
                image_features=torch.cat(all_image_features),
                text_features=torch.cat(all_text_features),
                logit_scale=logit_scale.cpu(),
            )
            loss = cumulative_loss / num_samples
            metrics.update(
                {**val_metrics, "clip_val_loss": loss.item(), "epoch": epoch, "num_samples": num_samples}
            )
            if gen_loss is not None:
                gen_loss = cumulative_gen_loss / num_samples
                metrics.update({"val_generative_loss": gen_loss.item()})

    if not metrics:
        return metrics

    logging.info(
        f"Eval Epoch: {epoch} "
        + "\t".join([f"{k}: {round(v, 4):.4f}" for k, v in metrics.items()])
    )

    log_data = {"val/" + name: val for name, val in metrics.items()}

    if args.save_logs:
        if tb_writer is not None:
            for name, val in log_data.items():
                tb_writer.add_scalar(name, val, epoch)

        with open(os.path.join(args.checkpoint_path, "results.jsonl"), "a+") as f:
            f.write(json.dumps(metrics))
            f.write("\n")

    if args.wandb:
        assert wandb is not None, 'Please install wandb.'
        if 'train' in data:
            dataloader = data['train'].dataloader
            num_batches_per_epoch = dataloader.num_batches // args.accum_freq
            step = num_batches_per_epoch * epoch
        else:
            step = None
        log_data['epoch'] = epoch
        wandb.log(log_data, step=step)

    return metrics


def get_clip_metrics(image_features, text_features, logit_scale):
    metrics = {}
    logits_per_image = (logit_scale * image_features @ text_features.t()).detach().cpu()
    logits_per_text = logits_per_image.t().detach().cpu()

    logits = {"image_to_text": logits_per_image, "text_to_image": logits_per_text}
    ground_truth = torch.arange(len(text_features)).view(-1, 1)

    for name, logit in logits.items():
        ranking = torch.argsort(logit, descending=True)
        preds = torch.where(ranking == ground_truth)[1]
        preds = preds.detach().cpu().numpy()
        metrics[f"{name}_mean_rank"] = preds.mean() + 1
        metrics[f"{name}_median_rank"] = np.floor(np.median(preds)) + 1
        for k in [1, 5, 10]:
            metrics[f"{name}_R@{k}"] = np.mean(preds < k)

    return metrics


def maybe_compute_generative_loss(model_out):
    if "logits" in model_out and "labels" in model_out:
        token_logits = model_out["logits"]
        token_labels = model_out["labels"]
        return F.cross_entropy(token_logits.permute(0, 2, 1), token_labels)