Speaker diarization example with onnxruntime Python API (#1395)

.github/workflows/speaker-diarization.yaml

@@ -0,0 +1,98 @@
+name: speaker-diarization
+
+on:
+  push:
+    branches:
+      - speaker-diarization
+  workflow_dispatch:
+
+concurrency:
+  group: speaker-diarization-${{ github.ref }}
+  cancel-in-progress: true
+
+jobs:
+  linux:
+    name: speaker diarization
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        os: [macos-latest]
+        python-version: ["3.10"]
+
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
+
+      - name: ccache
+        uses: hendrikmuhs/ccache-action@v1.2
+        with:
+          key: ${{ matrix.os }}-speaker-diarization
+
+      - name: Setup Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install pyannote
+        shell: bash
+        run: |
+          pip install pyannote.audio onnx onnxruntime
+
+      - name: Install sherpa-onnx from source
+        shell: bash
+        run: |
+          python3 -m pip install --upgrade pip
+          python3 -m pip install wheel twine setuptools
+
+          export CMAKE_CXX_COMPILER_LAUNCHER=ccache
+          export PATH="/usr/lib/ccache:/usr/local/opt/ccache/libexec:$PATH"
+
+          cat sherpa-onnx/python/sherpa_onnx/__init__.py
+
+          python3 setup.py bdist_wheel
+          ls -lh dist
+          pip install ./dist/*.whl
+
+      - name: Run tests
+        shell: bash
+        run: |
+          pushd scripts/pyannote/segmentation
+
+          python3 -c "import sherpa_onnx; print(sherpa_onnx.__file__)"
+          python3 -c "import sherpa_onnx; print(sherpa_onnx.__version__)"
+          python3 -c "import sherpa_onnx; print(dir(sherpa_onnx))"
+
+          curl -SL -O https://huggingface.co/csukuangfj/pyannote-models/resolve/main/segmentation-3.0/pytorch_model.bin
+
+          test_wavs=(
+            0-two-speakers-zh.wav
+            1-two-speakers-en.wav
+            2-two-speakers-en.wav
+            3-two-speakers-en.wav
+          )
+
+          for w in ${test_wavs[@]}; do
+            curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-segmentation-models/$w
+          done
+
+          soxi *.wav
+
+          curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-segmentation-models/sherpa-onnx-pyannote-segmentation-3-0.tar.bz2
+          tar xvf sherpa-onnx-pyannote-segmentation-3-0.tar.bz2
+          rm sherpa-onnx-pyannote-segmentation-3-0.tar.bz2
+          ls -lh sherpa-onnx-pyannote-segmentation-3-0
+
+          curl -SL -O https://github.com/k2-fsa/sherpa-onnx/releases/download/speaker-recongition-models/3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx
+
+          for w in ${test_wavs[@]}; do
+            echo "---------test $w (onnx)----------"
+            time ./speaker-diarization-onnx.py \
+              --seg-model ./sherpa-onnx-pyannote-segmentation-3-0/model.onnx \
+              --speaker-embedding-model ./3dspeaker_speech_eres2net_base_sv_zh-cn_3dspeaker_16k.onnx \
+              --wav $w
+
+            echo "---------test $w (torch)----------"
+            time ./speaker-diarization-torch.py  --wav $w
+          done

.gitignore

@@ -118,3 +118,5 @@ vits-melo-tts-zh_en
 *.o
 *.ppu
 sherpa-onnx-online-punct-en-2024-08-06
+*.mp4
+*.mp3

scripts/pyannote/segmentation/README.md

@@ -0,0 +1,44 @@
+# File description
+
+Please download test wave files from
+https://github.com/k2-fsa/sherpa-onnx/releases/tag/speaker-segmentation-models
+
+## 0-two-speakers-zh.wav
+
+This file is from
+https://www.modelscope.cn/models/iic/speech_campplus_speaker-diarization_common/file/view/master?fileName=examples%252F2speakers_example.wav&status=0
+
+Note that we have renamed it from `2speakers_example.wav` to `0-two-speakers-zh.wav`.
+
+## 1-two-speakers-en.wav
+
+This file is from
+https://github.com/pengzhendong/pyannote-onnx/blob/master/data/test_16k.wav
+and it contains speeches from two speakers.
+
+Note that we have renamed it from `test_16k.wav` to `1-two-speakers-en.wav`
+
+
+## 2-two-speakers-en.wav
+This file is from
+https://huggingface.co/spaces/Xenova/whisper-speaker-diarization
+
+Note that the original file is `./fcf059e3-689f-47ec-a000-bdace87f0113.mp4`.
+We use the following commands to convert it to `2-two-speakers-en.wav`.
+
+```bash
+ffmpeg -i ./fcf059e3-689f-47ec-a000-bdace87f0113.mp4 -ac 1 -ar 16000 ./2-two-speakers-en.wav
+```
+
+## 3-two-speakers-en.wav
+
+This file is from
+https://aws.amazon.com/blogs/machine-learning/deploy-a-hugging-face-pyannote-speaker-diarization-model-on-amazon-sagemaker-as-an-asynchronous-endpoint/
+
+Note that the original file is `ML16091-Audio.mp3`. We use the following
+commands to convert it to `3-two-speakers-en.wav`
+
+
+```bash
+sox ML16091-Audio.mp3 3-two-speakers-en.wav
+```

scripts/pyannote/segmentation/speaker-diarization-onnx.py

@@ -0,0 +1,488 @@
+#!/usr/bin/env python3
+# Copyright    2024  Xiaomi Corp.        (authors: Fangjun Kuang)
+
+"""
+Please refer to
+https://github.com/k2-fsa/sherpa-onnx/blob/master/.github/workflows/speaker-diarization.yaml
+for usages.
+"""
+
+import argparse
+from datetime import timedelta
+from pathlib import Path
+from typing import List
+
+import librosa
+import numpy as np
+import onnxruntime as ort
+import sherpa_onnx
+import soundfile as sf
+from numpy.lib.stride_tricks import as_strided
+
+
+class Segment:
+    def __init__(
+        self,
+        start,
+        end,
+        speaker,
+    ):
+        assert start < end
+        self.start = start
+        self.end = end
+        self.speaker = speaker
+
+    def merge(self, other, gap=0.5):
+        assert self.speaker == other.speaker, (self.speaker, other.speaker)
+        if self.end < other.start and self.end + gap >= other.start:
+            return Segment(start=self.start, end=other.end, speaker=self.speaker)
+        elif other.end < self.start and other.end + gap >= self.start:
+            return Segment(start=other.start, end=self.end, speaker=self.speaker)
+        else:
+            return None
+
+    @property
+    def duration(self):
+        return self.end - self.start
+
+    def __str__(self):
+        s = f"{timedelta(seconds=self.start)}"[:-3]
+        s += " --> "
+        s += f"{timedelta(seconds=self.end)}"[:-3]
+        s += f" speaker_{self.speaker:02d}"
+        return s
+
+
+def merge_segment_list(in_out: List[Segment], min_duration_off: float):
+    changed = True
+    while changed:
+        changed = False
+        for i in range(len(in_out)):
+            if i + 1 >= len(in_out):
+                continue
+
+            new_segment = in_out[i].merge(in_out[i + 1], gap=min_duration_off)
+            if new_segment is None:
+                continue
+            del in_out[i + 1]
+            in_out[i] = new_segment
+            changed = True
+            break
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--seg-model",
+        type=str,
+        required=True,
+        help="Path to model.onnx for segmentation",
+    )
+    parser.add_argument(
+        "--speaker-embedding-model",
+        type=str,
+        required=True,
+        help="Path to model.onnx for speaker embedding extractor",
+    )
+    parser.add_argument("--wav", type=str, required=True, help="Path to test.wav")
+
+    return parser.parse_args()
+
+
+class OnnxSegmentationModel:
+    def __init__(self, filename):
+        session_opts = ort.SessionOptions()
+        session_opts.inter_op_num_threads = 1
+        session_opts.intra_op_num_threads = 1
+
+        self.session_opts = session_opts
+
+        self.model = ort.InferenceSession(
+            filename,
+            sess_options=self.session_opts,
+            providers=["CPUExecutionProvider"],
+        )
+
+        meta = self.model.get_modelmeta().custom_metadata_map
+        print(meta)
+
+        self.window_size = int(meta["window_size"])
+        self.sample_rate = int(meta["sample_rate"])
+        self.window_shift = int(0.1 * self.window_size)
+        self.receptive_field_size = int(meta["receptive_field_size"])
+        self.receptive_field_shift = int(meta["receptive_field_shift"])
+        self.num_speakers = int(meta["num_speakers"])
+        self.powerset_max_classes = int(meta["powerset_max_classes"])
+        self.num_classes = int(meta["num_classes"])
+
+    def __call__(self, x):
+        """
+        Args:
+          x: (N, num_samples)
+        Returns:
+          A tensor of shape (N, num_frames, num_classes)
+        """
+        x = np.expand_dims(x, axis=1)
+
+        (y,) = self.model.run(
+            [self.model.get_outputs()[0].name], {self.model.get_inputs()[0].name: x}
+        )
+
+        return y
+
+
+def load_wav(filename, expected_sample_rate) -> np.ndarray:
+    audio, sample_rate = sf.read(filename, dtype="float32", always_2d=True)
+    audio = audio[:, 0]  # only use the first channel
+    if sample_rate != expected_sample_rate:
+        audio = librosa.resample(
+            audio,
+            orig_sr=sample_rate,
+            target_sr=expected_sample_rate,
+        )
+    return audio
+
+
+def get_powerset_mapping(num_classes, num_speakers, powerset_max_classes):
+    mapping = np.zeros((num_classes, num_speakers))
+
+    k = 1
+    for i in range(1, powerset_max_classes + 1):
+        if i == 1:
+            for j in range(0, num_speakers):
+                mapping[k, j] = 1
+                k += 1
+        elif i == 2:
+            for j in range(0, num_speakers):
+                for m in range(j + 1, num_speakers):
+                    mapping[k, j] = 1
+                    mapping[k, m] = 1
+                    k += 1
+        elif i == 3:
+            raise RuntimeError("Unsupported")
+
+    return mapping
+
+
+def to_multi_label(y, mapping):
+    """
+    Args:
+      y: (num_chunks, num_frames, num_classes)
+    Returns:
+      A tensor of shape (num_chunks, num_frames, num_speakers)
+    """
+    y = np.argmax(y, axis=-1)
+    labels = mapping[y.reshape(-1)].reshape(y.shape[0], y.shape[1], -1)
+    return labels
+
+
+# speaker count per frame
+def speaker_count(labels, seg_m):
+    """
+    Args:
+      labels: (num_chunks, num_frames, num_speakers)
+      seg_m: Segmentation model
+    Returns:
+      A integer array of shape (num_total_frames,)
+    """
+    labels = labels.sum(axis=-1)
+    # Now labels: (num_chunks, num_frames)
+
+    num_frames = (
+        int(
+            (seg_m.window_size + (labels.shape[0] - 1) * seg_m.window_shift)
+            / seg_m.receptive_field_shift
+        )
+        + 1
+    )
+    ans = np.zeros((num_frames,))
+    count = np.zeros((num_frames,))
+
+    for i in range(labels.shape[0]):
+        this_chunk = labels[i]
+        start = int(i * seg_m.window_shift / seg_m.receptive_field_shift + 0.5)
+        end = start + this_chunk.shape[0]
+        ans[start:end] += this_chunk
+        count[start:end] += 1
+
+    ans /= np.maximum(count, 1e-12)
+
+    return (ans + 0.5).astype(np.int8)
+
+
+def load_speaker_embedding_model(filename):
+    config = sherpa_onnx.SpeakerEmbeddingExtractorConfig(
+        model=filename,
+        num_threads=1,
+        debug=0,
+    )
+    if not config.validate():
+        raise ValueError(f"Invalid config. {config}")
+    extractor = sherpa_onnx.SpeakerEmbeddingExtractor(config)
+    return extractor
+
+
+def get_embeddings(embedding_filename, audio, labels, seg_m, exclude_overlap):
+    """
+    Args:
+      embedding_filename: Path to the speaker embedding extractor model
+      audio: (num_samples,)
+      labels: (num_chunks, num_frames, num_speakers)
+      seg_m: segmentation model
+    Returns:
+      Return (num_chunks, num_speakers, embedding_dim)
+    """
+    if exclude_overlap:
+        labels = labels * (labels.sum(axis=-1, keepdims=True) < 2)
+
+    extractor = load_speaker_embedding_model(embedding_filename)
+    buffer = np.empty(seg_m.window_size)
+    num_chunks, num_frames, num_speakers = labels.shape
+
+    ans_chunk_speaker_pair = []
+    ans_embeddings = []
+
+    for i in range(num_chunks):
+        labels_T = labels[i].T
+        # t: (num_speakers, num_frames)
+
+        sample_offset = i * seg_m.window_shift
+
+        for j in range(num_speakers):
+            frames = labels_T[j]
+            if frames.sum() < 10:
+                # skip segment less than 20 frames, i.e., about 0.2 seconds
+                continue
+
+            start = None
+            start_samples = 0
+            idx = 0
+            for k in range(num_frames):
+                if frames[k] != 0:
+                    if start is None:
+                        start = k
+                elif start is not None:
+                    start_samples = (
+                        int(start / num_frames * seg_m.window_size) + sample_offset
+                    )
+                    end_samples = (
+                        int(k / num_frames * seg_m.window_size) + sample_offset
+                    )
+                    num_samples = end_samples - start_samples
+                    buffer[idx : idx + num_samples] = audio[start_samples:end_samples]
+                    idx += num_samples
+
+                    start = None
+            if start is not None:
+                start_samples = (
+                    int(start / num_frames * seg_m.window_size) + sample_offset
+                )
+                end_samples = int(k / num_frames * seg_m.window_size) + sample_offset
+                num_samples = end_samples - start_samples
+                buffer[idx : idx + num_samples] = audio[start_samples:end_samples]
+                idx += num_samples
+
+            stream = extractor.create_stream()
+            stream.accept_waveform(sample_rate=seg_m.sample_rate, waveform=buffer[:idx])
+            stream.input_finished()
+
+            assert extractor.is_ready(stream)
+            embedding = extractor.compute(stream)
+            embedding = np.array(embedding)
+
+            ans_chunk_speaker_pair.append([i, j])
+            ans_embeddings.append(embedding)
+
+    assert len(ans_chunk_speaker_pair) == len(ans_embeddings), (
+        len(ans_chunk_speaker_pair),
+        len(ans_embeddings),
+    )
+    return ans_chunk_speaker_pair, np.array(ans_embeddings)
+
+
+def main():
+    args = get_args()
+    assert Path(args.seg_model).is_file(), args.seg_model
+    assert Path(args.wav).is_file(), args.wav
+
+    seg_m = OnnxSegmentationModel(args.seg_model)
+    audio = load_wav(args.wav, seg_m.sample_rate)
+    # audio: (num_samples,)
+
+    num = (audio.shape[0] - seg_m.window_size) // seg_m.window_shift + 1
+
+    samples = as_strided(
+        audio,
+        shape=(num, seg_m.window_size),
+        strides=(seg_m.window_shift * audio.strides[0], audio.strides[0]),
+    )
+
+    # or use torch.Tensor.unfold
+    #  samples = torch.from_numpy(audio).unfold(0, seg_m.window_size, seg_m.window_shift).numpy()
+
+    if (
+        audio.shape[0] < seg_m.window_size
+        or (audio.shape[0] - seg_m.window_size) % seg_m.window_shift > 0
+    ):
+        has_last_chunk = True
+    else:
+        has_last_chunk = False
+
+    num_chunks = samples.shape[0]
+    batch_size = 32
+    output = []
+    for i in range(0, num_chunks, batch_size):
+        start = i
+        end = i + batch_size
+        # it's perfectly ok to use end > num_chunks
+        y = seg_m(samples[start:end])
+        output.append(y)
+
+    if has_last_chunk:
+        last_chunk = audio[num_chunks * seg_m.window_shift :]  # noqa
+        pad_size = seg_m.window_size - last_chunk.shape[0]
+        last_chunk = np.pad(last_chunk, (0, pad_size))
+        last_chunk = np.expand_dims(last_chunk, axis=0)
+        y = seg_m(last_chunk)
+        output.append(y)
+
+    y = np.vstack(output)
+    # y: (num_chunks, num_frames, num_classes)
+
+    mapping = get_powerset_mapping(
+        num_classes=seg_m.num_classes,
+        num_speakers=seg_m.num_speakers,
+        powerset_max_classes=seg_m.powerset_max_classes,
+    )
+    labels = to_multi_label(y, mapping=mapping)
+    # labels: (num_chunks, num_frames, num_speakers)
+
+    inactive = (labels.sum(axis=1) == 0).astype(np.int8)
+    # inactive: (num_chunks, num_speakers)
+
+    speakers_per_frame = speaker_count(labels=labels, seg_m=seg_m)
+    # speakers_per_frame: (num_frames, speakers_per_frame)
+
+    if speakers_per_frame.max() == 0:
+        print("No speakers found in the audio file!")
+        return
+
+    # if users specify only 1 speaker for clustering, then return the
+    # result directly
+
+    # Now, get embeddings
+    chunk_speaker_pair, embeddings = get_embeddings(
+        args.speaker_embedding_model,
+        audio=audio,
+        labels=labels,
+        seg_m=seg_m,
+        #  exclude_overlap=True,
+        exclude_overlap=False,
+    )
+    # chunk_speaker_pair: a list of (chunk_idx, speaker_idx)
+    # embeddings: (batch_size, embedding_dim)
+
+    # Please change num_clusters or threshold by yourself.
+    clustering_config = sherpa_onnx.FastClusteringConfig(num_clusters=2)
+    #  clustering_config = sherpa_onnx.FastClusteringConfig(threshold=0.8)
+    clustering = sherpa_onnx.FastClustering(clustering_config)
+    cluster_labels = clustering(embeddings)
+
+    chunk_speaker_to_cluster = dict()
+    for (chunk_idx, speaker_idx), cluster_idx in zip(
+        chunk_speaker_pair, cluster_labels
+    ):
+        if inactive[chunk_idx, speaker_idx] == 1:
+            print("skip ", chunk_idx, speaker_idx)
+            continue
+        chunk_speaker_to_cluster[(chunk_idx, speaker_idx)] = cluster_idx
+
+    num_speakers = max(cluster_labels) + 1
+    relabels = np.zeros((labels.shape[0], labels.shape[1], num_speakers))
+    for i in range(labels.shape[0]):
+        for j in range(labels.shape[1]):
+            for k in range(labels.shape[2]):
+                if (i, k) not in chunk_speaker_to_cluster:
+                    continue
+                t = chunk_speaker_to_cluster[(i, k)]
+
+                if labels[i, j, k] == 1:
+                    relabels[i, j, t] = 1
+
+    num_frames = (
+        int(
+            (seg_m.window_size + (relabels.shape[0] - 1) * seg_m.window_shift)
+            / seg_m.receptive_field_shift
+        )
+        + 1
+    )
+
+    count = np.zeros((num_frames, relabels.shape[-1]))
+    for i in range(relabels.shape[0]):
+        this_chunk = relabels[i]
+        start = int(i * seg_m.window_shift / seg_m.receptive_field_shift + 0.5)
+        end = start + this_chunk.shape[0]
+        count[start:end] += this_chunk
+
+    if has_last_chunk:
+        stop_frame = int(audio.shape[0] / seg_m.receptive_field_shift)
+        count = count[:stop_frame]
+
+    sorted_count = np.argsort(-count, axis=-1)
+    final = np.zeros((count.shape[0], count.shape[1]))
+
+    for i, (c, sc) in enumerate(zip(speakers_per_frame, sorted_count)):
+        for k in range(c):
+            final[i, sc[k]] = 1
+
+    min_duration_off = 0.5
+    min_duration_on = 0.3
+    onset = 0.5
+    offset = 0.5
+    # final: (num_frames, num_speakers)
+
+    final = final.T
+    for kk in range(final.shape[0]):
+        segment_list = []
+        frames = final[kk]
+
+        is_active = frames[0] > onset
+
+        start = None
+        if is_active:
+            start = 0
+        scale = seg_m.receptive_field_shift / seg_m.sample_rate
+        scale_offset = seg_m.receptive_field_size / seg_m.sample_rate * 0.5
+        for i in range(1, len(frames)):
+            if is_active:
+                if frames[i] < offset:
+                    segment = Segment(
+                        start=start * scale + scale_offset,
+                        end=i * scale + scale_offset,
+                        speaker=kk,
+                    )
+                    segment_list.append(segment)
+                    is_active = False
+            else:
+                if frames[i] > onset:
+                    start = i
+                    is_active = True
+
+        if is_active:
+            segment = Segment(
+                start=start * scale + scale_offset,
+                end=(len(frames) - 1) * scale + scale_offset,
+                speaker=kk,
+            )
+            segment_list.append(segment)
+
+        if len(segment_list) > 1:
+            merge_segment_list(segment_list, min_duration_off=min_duration_off)
+            for s in segment_list:
+                if s.duration < min_duration_on:
+                    continue
+                print(s)
+
+
+if __name__ == "__main__":
+    main()

scripts/pyannote/segmentation/speaker-diarization-torch.py

@@ -0,0 +1,86 @@
+#!/usr/bin/env python3
+
+"""
+Please refer to
+https://github.com/k2-fsa/sherpa-onnx/blob/master/.github/workflows/speaker-diarization.yaml
+for usages.
+"""
+
+"""
+1. Go to https://huggingface.co/hbredin/wespeaker-voxceleb-resnet34-LM/tree/main
+wget https://huggingface.co/hbredin/wespeaker-voxceleb-resnet34-LM/resolve/main/speaker-embedding.onnx
+
+2. Change line 166 of pyannote/audio/pipelines/speaker_diarization.py
+
+```
+            #  self._embedding = PretrainedSpeakerEmbedding(
+            #      self.embedding, use_auth_token=use_auth_token
+            #  )
+            self._embedding = embedding
+```
+"""
+
+import argparse
+from pathlib import Path
+
+import torch
+from pyannote.audio import Model
+from pyannote.audio.pipelines import SpeakerDiarization as SpeakerDiarizationPipeline
+from pyannote.audio.pipelines.speaker_verification import (
+    ONNXWeSpeakerPretrainedSpeakerEmbedding,
+)
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--wav", type=str, required=True, help="Path to test.wav")
+
+    return parser.parse_args()
+
+
+def build_pipeline():
+    embedding_filename = "./speaker-embedding.onnx"
+    if Path(embedding_filename).is_file():
+        # You need to modify line 166
+        # of pyannote/audio/pipelines/speaker_diarization.py
+        # Please see the comments at the start of this script for details
+        embedding = ONNXWeSpeakerPretrainedSpeakerEmbedding(embedding_filename)
+    else:
+        embedding = "hbredin/wespeaker-voxceleb-resnet34-LM"
+
+    pt_filename = "./pytorch_model.bin"
+    segmentation = Model.from_pretrained(pt_filename)
+    segmentation.eval()
+
+    pipeline = SpeakerDiarizationPipeline(
+        segmentation=segmentation,
+        embedding=embedding,
+        embedding_exclude_overlap=True,
+    )
+
+    params = {
+        "clustering": {
+            "method": "centroid",
+            "min_cluster_size": 12,
+            "threshold": 0.7045654963945799,
+        },
+        "segmentation": {"min_duration_off": 0.5},
+    }
+
+    pipeline.instantiate(params)
+    return pipeline
+
+
+@torch.no_grad()
+def main():
+    args = get_args()
+    assert Path(args.wav).is_file(), args.wav
+    pipeline = build_pipeline()
+    print(pipeline)
+    t = pipeline(args.wav)
+    print(type(t))
+    print(t)
+
+
+if __name__ == "__main__":
+    main()

sherpa-onnx/csrc/fast-clustering.cc

@@ -52,7 +52,7 @@ class FastClustering::Impl {
     std::vector<double> height(num_rows - 1);
 
     fastclustercpp::hclust_fast(num_rows, distance.data(),
-                                fastclustercpp::HCLUST_METHOD_SINGLE,
+                                fastclustercpp::HCLUST_METHOD_COMPLETE,
                                 merge.data(), height.data());
 
     std::vector<int32_t> labels(num_rows);