INNER CODE UNIT · Python

generate_embeddings_clip

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

def generate_embeddings_clip(
    model: torch.nn.Module, data_loader: DataLoader, device: torch.device
) -> Tuple[torch.Tensor, torch.Tensor, float]:
    model.eval()
    embedding_all = []
    labels_all = []
    elapsed = 0.0

    with torch.inference_mode():
        for inputs, labels in tqdm.tqdm(data_loader):
            start_time = time.monotonic()
            outputs = model(**{"pixel_values": inputs.to(device)})
            elapsed += time.monotonic() - start_time
            embedding_all.append(outputs.pooler_output)
            labels_all.append(labels)

        return torch.cat(embedding_all), torch.cat(labels_all), elapsed

View source record →