p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/train_dynamics.py (python)
def _calculate_top_k_accuracy(
    token_logits: jax.Array,
    video_tokens: jax.Array,
    mask: jax.Array,
    k: int,
) -> jax.Array:
    topk_indices = jnp.argsort(token_logits, axis=-1)[..., -k:]
    topk_correct = jnp.any(topk_indices == video_tokens[..., None], axis=-1)
    topk_acc = (mask * topk_correct).sum() / mask.sum()
    return topk_acc


def _calculate_step_metrics(
    outputs: dict[str, jax.Array],
    gt: jax.Array,
    num_actions: int,
    num_patch_latents: int,
) -> tuple[jax.Array, dict]:
    mask = outputs["mask"]
    outputs["token_logits"] = outputs["token_logits"].astype(jnp.float32)
    ce_loss = optax.softmax_cross_entropy_with_integer_labels(
        outputs["token_logits"], outputs["video_tokens"]
    )
    ce_loss = (mask * ce_loss).sum() / mask.sum()

    masked_token_top_1_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask, 1
    )
    masked_token_top_2_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask, 2
    )
    masked_token_top_5_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask, 5
    )
    masked_token_top_16_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask, 16
    )

    select_probs = jax.nn.softmax(outputs["token_logits"])
    gt_val = gt.clip(0, 1).reshape(-1, *gt.shape[2:])
    recon = outputs["recon"].clip(0, 1).reshape(-1, *outputs["recon"].shape[2:])
    psnr = jnp.asarray(pix.psnr(gt_val, recon)).mean()
    ssim = jnp.asarray(pix.ssim(gt_val, recon)).mean()
    _, index_counts_tokenizer = jnp.unique_counts(
        jnp.ravel(outputs["video_tokens"]),
        size=num_patch_latents,
        fill_value=0,
    )
    codebook_usage_tokenizer = (index_counts_tokenizer != 0).mean()
    metrics = dict(
        cross_entropy_loss=ce_loss,
        masked_token_top1_accuracy=masked_token_top_1_acc,
        masked_token_top2_accuracy=masked_token_top_2_acc,
        masked_token_top5_accuracy=masked_token_top_5_acc,
        masked_token_top16_accuracy=masked_token_top_16_acc,