# Get the embeddings for the dataset

from pyannote.audio import Model

from pyannote.audio import Inference

model = Model.from_pretrained("pyannote/embedding",

use_auth_token="<YOUR_HF_AUTH_TOKEN>")

inference = Inference(model, window="whole")

embeddings = []

ids = []

for clip_file in dataset:

embeddings.append(inference(clip_file))

ids.append(clip_file)

# Add the embeddings to chroma

collection = client.get_or_create_collection('voxceleb')

collection.add(ids=ids, embeddings=embeddings)

# Query with a test clip

test_embedding = inference('test.wav')

results = collection.query(query_embeddings=test_embedding)