Project: 068_Two_Effects_One_Trigger_On_the_Modality_Gap_Object | Function: compute_rmg (iclr)

system: You are an expert Python programmer implementing algorithms from scientific papers.
Complete the masked function based on the algorithm description.
Follow Python best practices and maintain consistent code style.
Return **only** the complete Python function code.

user: Complete this Python function:

**DESCRIPTION:** Computes the Relative Modality Gap (RMG) metric to measure modality separation while accounting for matching image-text pairs and effectively used embedding space, addressing limitations of the L2M metric

**KEY TERMS:** {
  "image_features_matching": "d(x_i, y_i)",
  "mean_img_similarity": "intra-modality distances for images",
  "mean_txt_similarity": "intra-modality distances for text",
  "normalizer": "denominator of RMG formula"
}

**CONTEXT CODE:**     #  torch.argsort runs out of memory for me (6GB VRAM) so I move to CPU for sorting
    # dist_matrix = dist_matrix.cpu()
    # Sort in descending order; first is the biggest logit
    inds = torch.argsort(dist_matrix, dim=1, descending=True)
    text_to_image_recall = []
    for k in k_vals:
        # Extract top k indices only
        topk = inds[:, :k]
        # Correct iff one of the top_k values equals the correct image (as given by text_to_image_map)
        correct = torch.eq(topk, text_to_image_map.unsqueeze(-1)).any(dim=1)
        num_correct = correct.sum().item()
        text_to_image_recall.append(num_correct / num_text)

    # image-to-text recall
    dist_matrix = dist_matrix.T  # dist_matrix[i] gives logits for the ith image
    # Sort in descending order; first is the biggest logit
    inds = torch.argsort(dist_matrix, dim=1, descending=True)
    image_to_text_recall = []
    for k in k_vals:
        # Extract top k indices only
        topk = inds[:, :k]
        correct = torch.zeros((num_im,), dtype=torch.bool).to(device)
        #  For each image, check whether one of the 5 relevant captions was retrieved
        # Check if image matches its ith caption (for i=0..4)
        for i in range(captions_per_image):
            contains_index = torch.eq(topk, image_to_text_map[:, i].unsqueeze(-1)).any(dim=1)
            correct = torch.logical_or(correct, contains_index)
        num_correct = correct.sum().item()
        image_to_text_recall.append(num_correct / num_im)

    return text_to_image_recall, image_to_text_recall

# Modality gap metrics
@torch.no_grad()
def compute_l2m(image_features: torch.Tensor, text_features: torch.Tensor) -> float:
    return (image_features.mean(axis=0) - text_features.mean(axis=0)).norm().item()

@torch.no_grad()
def compute_l2i(image_features: torch.Tensor, text_features: torch.Tensor, dataset, labels=None, text_to_image_map=None) -> float:
    if dataset == "imagenet":
        image_features_per_text = image_features
        text_feature_per_image = torch.stack([text_features[l] for l in labels], dim=0)
    elif dataset == "coco":
        image_features_per_text = torch.stack([image_features[i] for i in text_to_image_map], dim=0)
        text_feature_per_image = text_features
    else:
        raise NotImplementedError
    return (image_features_per_text - text_feature_per_image).norm(dim=-1).mean().item()

@torch.no_grad()
def compute_rmg(image_features, text_features, dataset, labels=None, text_to_image_map=None, image_to_text_map=None):
    raise NotImplementedError("Function implementation has been masked for testing")
    
@torch.no_grad()
def encode_clip(model, loader, tokenizer, prompts=None):
    device = next(model.parameters()).device

    with torch.amp.autocast(device_type=device.type):
        all_image_features, labels, all_text_features = [], [], []
        image_to_text_map, text_to_image_map = [], []
        captions_per_image = []
        text_index, image_index = 0, 0
        for batch in tqdm(loader, leave=False, total=len(loader)):
            if len(batch) == 2:
                inputs, label = batch
            else:
                raise NotImplementedError
            if isinstance(inputs, tuple):
                inputs = torch.stack(inputs, dim=0)
            
            image_features = model.encode_image(inputs.to(device))
            image_features /= image_features.norm(dim=-1, keepdim=True)
            all_image_features.append(image_features.cpu())

            if prompts is None: # coco
                text = label # shape B x 5 x 77
                batch_size, captions_per_image, _ = text.shape

                # Update text_to_image_map and image_to_text_map for this batch
                for i in range(batch_size):
                    # the next image corresponds to text captions [text_index ... text_index + captions_per_image - 1]
                    text_indices = list(range(text_index, text_index + captions_per_image))
                    image_to_text_map.append(text_indices)
                    text_index += captions_per_image

                    # Each of the next captions_per_image text captions correspond to the same image
                    text_to_image_map += [image_index] * captions_per_image
                    image_index += 1

                # B x 5 x 77 -> (B*5) x 77
                text = torch.flatten(text, start_dim=0, end_dim=1)
                try:
                    text_features = model.encode_text(text.to(device))
                except Exception:
                    text_features = []
                    for t in text:
                        text_features.append(model.encode_text(t.unsqueeze(0).to(device)).cpu())
                    text_features = torch.cat(text_features, dim=0)
                text_features /= text_features.norm(dim=-1, keepdim=True)
                all_text_features.append(text_features.cpu())
            else: # imagenet
                labels.append(label)
        
        all_image_features = torch.cat(all_image_features, dim=0)

        if prompts is not None: # imagenet
            labels = torch.cat(labels, dim=0)
            text = tokenizer(prompts).to(device)
            try:
                text_features = model.encode_text(text.to(device)).cpu()
            except Exception:
                text_features = []
                for t in text:
                    text_features.append(model.encode_text(t.unsqueeze(0)).cpu())
                text_features = torch.cat(text_features, dim=0)
            text_features /= text_features.norm(dim=-1, keepdim=True)
            return all_image_features, text_features, labels
        else:
            text_features = torch.cat(all_text_features, dim=0)
            text_to_image_map = torch.LongTensor(text_to_image_map).to(device)
            image_to_text_map = torch.LongTensor(image_to_text_map).to(device)
            return all_image_features, text_features, text_to_image_map, image_to_text_map

def main(args):
    args.output_dir = Path(args.output_dir) / "performance_vs_modality_gap"
    os.makedirs(args.output_dir, exist_ok=True)
    
    model_name, pretrain_dataset = args.model.split("__")
    model, _, preprocess = open_clip.create_model_and_transforms(model_name, pretrained=pretrain_dataset, cache_dir=args.cache_dir)
    tokenizer = open_clip.get_tokenizer(model_name)
    model = model.to(args.device)
    model.eval()
    tokenizer = partial(tokenizer, context_length=model.context_length)

    datasets = ["imagenet", "coco"] if args.dataset == "all" else [args.dataset]

    for dataset in datasets:
        print(5*"#", dataset, 5*"#")
        if dataset == "imagenet":
            data = ImageNet(IMAGENET_DIR, transform=preprocess)
            prompts = [f"A photo of a {obj}." for obj in openai_imagenet_classes]
            loader = DataLoader(
                dataset=data,
                batch_size=args.batch_size,
                shuffle=False,
                num_workers=8,
                collate_fn=None,
            )
        elif dataset == "coco":
            def prepend(text: list) -> list:


Complete the function implementation replacing "raise NotImplementedError".