p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] jasmine/train_dynamics.py (python)
def _calculate_top_k_accuracy(
    token_logits_BTNV: jax.Array,
    video_tokens_BTN: jax.Array,
    mask_BTN: jax.Array,
    k: int,
) -> jax.Array:
    _, topk_indices_BTNK = jax.lax.top_k(token_logits_BTNV, k)
    topk_correct = jnp.any(
        topk_indices_BTNK == video_tokens_BTN[..., jnp.newaxis], axis=-1
    )
    topk_acc = (mask_BTN * topk_correct).sum() / mask_BTN.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_BTN = 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_BTN * ce_loss).sum() / mask_BTN.sum()

    masked_token_top_1_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 1
    )
    masked_token_top_2_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 2
    )
    masked_token_top_5_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 5
    )
    masked_token_top_16_acc = _calculate_top_k_accuracy(
        outputs["token_logits"], outputs["video_tokens"], mask_BTN, 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,
[content] jasmine/train_dynamics.py (python)
def _calculate_top_k_accuracy(
    token_logits_BTNV: jax.Array,
    video_tokens_BTN: jax.Array,
    mask_BTN: jax.Array,
    k: int,
) -> jax.Array:
    _, topk_indices_BTNK = jax.lax.top_k(token_logits_BTNV, k)
    topk_correct = jnp.any(
        topk_indices_BTNK == video_tokens_BTN[..., jnp.newaxis], axis=-1
    )
    topk_acc = (mask_BTN * topk_correct).sum() / mask_BTN.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_BTN = 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"]
    )
<<<<<<< HEAD
    ce_loss = (mask * ce_loss).sum() / mask.sum()
    acc = outputs["token_logits"].argmax(-1) == outputs["video_tokens"]
    acc = (mask * acc).sum() / mask.sum()
    z_val = jax.nn.logsumexp(outputs["token_logits"], axis=-1)
    z_loss_metric = (mask * (z_val**2)).sum() / mask.sum()
=======
    ce_loss = (mask_BTN * ce_loss).sum() / mask_BTN.sum()

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

>>>>>>> main
    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,