INNER CODE UNIT · Python

run_clip_model

openfoodfacts/openfoodfacts-ai · logo-ann/benchmarks/embedding_models_benchmark/main.py:152

def run_clip_model(
    root_dir: Path,
    split_set: Optional[Set[str]],
    model_name: str,
    batch_size: int,
    device: torch.device,
):
    model = CLIPModel.from_pretrained(f"openai/{model_name}").vision_model
    model.to(device)
    processor = CLIPProcessor.from_pretrained(f"openai/{model_name}")
    transform_func = lambda x: processor(images=x, return_tensors="pt")["pixel_values"][
        0
    ]
    dataset = LogoDataset(root_dir, transform_func, split_set)
    data_loader = DataLoader(dataset, batch_size, num_workers=2)
    embeddings, labels, elapsed = generate_embeddings_clip(model, data_loader, device)
    return embeddings, labels, dataset, elapsed

View source record →