Given the code below, provide modifications and discussion to consider how your code would
scale to thousands or even millions of images or models:

import torch
import clip
import time
from memory_profiler import memory_usage
from torch.utils.data import Dataset, DataLoader
from utils import load_data, ImageCaptionDataset

def compute_similarities(dataloader, model, device):
"""
Compute similarity between image and caption embeddings using CLIP.
"""
similarities = []

for i, (images, texts) in enumerate(dataloader):
# print(images.shape, texts.shape)
images = images.to(device)
texts = texts.to(device)

with torch.no_grad():
    image_features = model.encode_image(images)
    text_features = model.encode_text(texts)
    image_features /= image_features.norm(dim=-1, keepdim=True)
    text_features /= text_features.norm(dim=-1, keepdim=True)
    similarity = (image_features * text_features).sum(dim=-1).cpu().numpy()
    similarities.extend(similarity.tolist())

print(similarities)

return similarities

def filter_high_quality_pairs(df, threshold=0.3):
"""
Filter out low-quality image-caption pairs and save results to CSV.
"""

Filter out pairs with similarity below the threshold

high_quality_df = df[df['similarity'] >= threshold]

Save the filtered DataFrame back to CSV

high_quality_df.to_csv('filtered_results.csv', index=False)
print(f"Filtered results saved to 'filtered_results.csv' with {len(high_quality_df)} high-quality pairs.")

def main():
import argparse
parser = argparse.ArgumentParser(description="Compute image-text similarity using CLIP.")
parser.add_argument('--batch-size', type=int, default=4,
help='Batch size to use')
parser.add_argument('--csv', type=str, required=True, help="Path to the CSV file.")
parser.add_argument('--images_dir', type=str, required=True, help="Directory containing images.")

parser.add_argument('--csv', type=str, default='/content/drive/MyDrive/challenge_set/challenge_set.csv',
help="Path to the CSV file.")
parser.add_argument('--images_dir', type=str, default='/content/drive/MyDrive/challenge_set',
help="Directory containing images.")

parser.add_argument('--output_csv', type=str, default='output.csv', help="Output CSV file with similarity scores.")
parser.add_argument('--metrics_file', type=str, default='metrics.txt', help="File to store performance metrics.")
parser.add_argument('--device', type=str, default=None,
help='Device to use. Like cuda, cuda:0 or cpu')
parser.add_argument('--clip-model', type=str, default='ViT-B/32',
help='CLIP model to use')

args = parser.parse_args()

start_time = time.time()
mem_usage_start = memory_usage()[0]

if args.device is None:
device = torch.device('cuda' if (torch.cuda.is_available()) else 'cpu')
else:
device = torch.device(args.device)

Load CLIP model

model, preprocess = clip.load(args.clip_model, device=device)

Load data

df = load_data(args.csv, args.images_dir)
dataset = ImageCaptionDataset(df, transform=preprocess, tokenizer=clip.tokenize)
dataloader = DataLoader(dataset, batch_size=args.batch_size, shuffle=False)

Compute similarity

similarities = compute_similarities(dataloader, model, device)
df['similarity'] = similarities

Save updated CSV

df.to_csv(args.output_csv, index=False)
print(f"Similarity scores saved to {args.output_csv}")

filter_high_quality_pairs(df, threshold=0.3)

end_time = time.time()
mem_usage_end = memory_usage()[0]

Record metrics

total_time = end_time - start_time
total_memory = mem_usage_end - mem_usage_start
with open(args.metrics_file, 'w') as f:
f.write(f"Total Execution Time: {total_time:.2f} seconds\n")
f.write(f"Total Memory Usage: {total_memory:.2f} MiB\n")
print(f"Performance metrics saved to {args.metrics_file}")

if name == "main":
main()