def update_weight_sheared_llama(
    prob_map: dict[str, float], ref_loss: dict[str, float], curr_loss: dict[str, float]
) -> dict[str, float]:
    """
    Args:
        prob_map: dataset name -> prob
        ref_loss: dataset name -> ref loss
        curr_loss: dataset name -> curr loss

    Returns:
        prob_map: updated prob map

    References:
        Dynamic Batch Loading in ShearedLlama (http://arxiv.org/abs/2310.06694)
    """
    task_types = [k for k in prob_map]
    original_weight = torch.tensor([prob_map[k] for k in task_types])
    diff = torch.tensor([curr_loss[k] - ref_loss[k] for k in task_types])
    eta = 1.0
    c = 1e-4

    updated_alpha = torch.log(original_weight) + eta * diff
    updated_alpha = torch.nn.functional.softmax(updated_alpha, dim=0)
    updated_domain_weights = (1 - c) * updated_alpha + c / len(task_types)
    updated_domain_weights = updated_domain_weights.detach().numpy().astype("float64")
    updated_domain_weights = updated_domain_weights / updated_domain_weights.sum()
    updated_domain_weights = updated_domain_weights.tolist()

    return {k: v for k, v in zip(task_types, updated_domain_weights)}

    Args:
        prob_map: dataset name -> prob
        ref_loss: dataset name -> ref loss
        curr_loss: dataset name -> curr loss

    Returns:
        prob_map: updated prob map

    References:
        Dynamic Batch Loading in ShearedLlama (http://arxiv.org/abs/2310.06694)
    
smoe/data/dynamic_selection.py