README.md

---
library_name: transformers
tags:
- prosody
- segmentation
- audio
- speech
language:
- sl
base_model:
- facebook/w2v-bert-2.0
---

# Wav2Vec2Bert Audio frame classifier for prosodic unit detection

This model predicts prosodic units on speech. For each 20ms frame the model
predicts 1 or 0, indicating whether there is a prosodic unit in this frame or
not.

This frame-level output can be grouped into events with the frames_to_intervals
function provided in the code snippets below.

It is known that the model is unreliable if the audio starts or ends within a
prosodic unit. This can be somewhat circumvented by 1) using the largest
possible chunks that will fit your machine and 2) use overlapping chunks and
combining results smartly.




## Model Details

### Model Description



- **Developed by:** Peter Rupnik, Nikola Ljubešić, Darinka Verdonik, Simona
  Majhenič
- **Funded by:** MEZZANINE project
- **Model type:** Wav2Vec2Bert for Audio Frame Classification
- **Language(s) (NLP):** Trained and tested on Slovenian, ATM unclear if usable
  cross-lingually
- **Finetuned from model:** facebook/w2v-bert-2.0

The model was trained on [ROG-Art dataset](http://hdl.handle.net/11356/1992), on
train split only.

### Model performance

We evaluate the model indirectly, and only care about the positive class:

1. first prosodic units (intervals with start and end times, e.g. `[0.123,
   5.546]`) are extracted from data and model outputs
2. if a predicted prosodic unit has an overlapping counterpart in true prosodic
   units, we count it as a True Positive. If there is no overlapping true
   counterpart, we count it as a False Positive, and if we have a true prosodic
   unit without a counterpart in predictions, we count that as a False Negative.
3. Based on the TP, FN, FP numbers recall, precision, and F1 score is
   calculated.

In this fashion we obtain the following metrics:

* Precision: 0.9423
* Recall: 0.7802
* F_1 score: 0.8538

![A gif illustrating correspondance between true and predicted prosodic
units](output.gif)

As seen in the gif image above, we observe generally good correspondence between true (blue) and predicted (orange) prosodic units, but there are cases where the grouping is incorrect: the model will annotate only a single prosodic unit where a human annotator would annotate two or more.

### Known limitations

* Edge cases: if the input audio starts or ends within a prosodic unit, there is a high  chance of not detecting the ending or starting prosodic unit.
* Unknown behaviour on non-speech audio: as of the time of writing, no tests were performed to check what happens in cases of music, noise, pure sine, ...
## Uses

### Simple use (short files)

For shorter audios that fit on your GPU the classifier can be used directly.
```python
import numpy as np

from datasets import Audio, Dataset
from transformers import AutoFeatureExtractor, Wav2Vec2BertForAudioFrameClassification
import torch
import numpy as np

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

model_name = "5roop/Wav2Vec2BertProsodicUnitsFrameClassifier"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
model = Wav2Vec2BertForAudioFrameClassification.from_pretrained(model_name).to(device)
f = "data/Rog-Art-N-G6007-P600702_181.070_211.070.wav"


def frames_to_intervals(frames: list) -> list[tuple]:
    from itertools import pairwise
    import pandas as pd

    results = []
    ndf = pd.DataFrame(
        data={
            "time_s": [0.020 * i for i in range(len(frames))],
            "frames": frames,
        }
    )
    ndf = ndf.dropna()
    indices_of_change = ndf.frames.diff()[ndf.frames.diff() != 0].index.values
    for si, ei in pairwise(indices_of_change):
        if ndf.loc[si : ei - 1, "frames"].mode()[0] == 0:
            pass
        else:
            results.append(
                (round(ndf.loc[si, "time_s"], 3), round(ndf.loc[ei - 1, "time_s"], 3))
            )
    return results


def evaluator(chunks):
    sampling_rate = chunks["audio"][0]["sampling_rate"]
    with torch.no_grad():
        inputs = feature_extractor(
            [i["array"] for i in chunks["audio"]],
            return_tensors="pt",
            sampling_rate=sampling_rate,
        ).to(device)
        logits = model(**inputs).logits
    y_pred_raw = np.array(logits.cpu())
    y_pred = y_pred_raw.argmax(axis=-1)
    prosodic_units = [frames_to_intervals(i) for i in y_pred]
    return {
        "y_pred": y_pred,
        "y_pred_logits": y_pred_raw,
        "prosodic_units": prosodic_units,
    }

# Create a dataset with a single instance and map our evaluator function on it:
ds = Dataset.from_dict({"audio": [f]}).cast_column("audio", Audio(16000, mono=True))
ds = ds.map(evaluator, batched=True, batch_size=1) # Adjust batch size according to your hardware specs
print(ds["y_pred"][0])
# Outputs: [0, 0, 1, 1, 1, 1, 1, ...]
print(ds["y_pred_logits"][0])
# Outputs:
# [[ 0.89419061, -0.77746612],
#  [ 0.44213724, -0.34862748],
#  [-0.08605709,  0.13012762],
# ....
print(ds["prosodic_units"][0])
# Outputs: [[0.04, 2.4], [3.52, 6.6], ....
```


### Inference on longer files
If the file is too big for straight-forward inference, some chunking needs to be
performed in order to process it. We know that for starts and ends of chunks the
probability of false negatives increases, so it is best to process the file with
some overlap between chunks or split it on silence. We illustrate the former
approach here:
```python
import numpy as np

from datasets import Audio, Dataset
from transformers import AutoFeatureExtractor, Wav2Vec2BertForAudioFrameClassification
import torch
import numpy as np

if torch.cuda.is_available():
    device = torch.device("cuda")
else:
    device = torch.device("cpu")

model_name = "5roop/Wav2Vec2BertProsodicUnitsFrameClassifier"
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
model = Wav2Vec2BertForAudioFrameClassification.from_pretrained(model_name).to(device)
f = "ROG/ROG-Art/WAV/Rog-Art-N-G5025-P600022.wav"

OVERLAP_S = 10
CHUNK_LENGTH_S = 30
SAMPLING_RATE = 16_000
OVERLAP_SAMPLES = OVERLAP_S * SAMPLING_RATE
CHUNK_LENGTH_SAMPLES = CHUNK_LENGTH_S * SAMPLING_RATE


def frames_to_intervals(frames: list) -> list[tuple]:
    from itertools import pairwise
    import pandas as pd

    results = []
    ndf = pd.DataFrame(
        data={
            "time_s": [0.020 * i for i in range(len(frames))],
            "frames": frames,
        }
    )
    ndf = ndf.dropna()
    indices_of_change = ndf.frames.diff()[ndf.frames.diff() != 0].index.values
    for si, ei in pairwise(indices_of_change):
        if ndf.loc[si : ei - 1, "frames"].mode()[0] == 0:
            pass
        else:
            results.append(
                (round(ndf.loc[si, "time_s"], 3), round(ndf.loc[ei - 1, "time_s"], 3))
            )
    return results


def merge_events(events: list[list[float]], centroids):
    flattened_events = []
    flattened_centroids = []
    for batch_idx, batch in enumerate(events):
        for event in batch:
            flattened_events.append(event)
            flattened_centroids.append(centroids[batch_idx])
    flattened_events.sort(key=lambda x: x[0])

    # Merged list to store final intervals
    merged = []

    for event, centroid in zip(flattened_events, flattened_centroids):
        if not merged:
            # If merged is empty, simply add the first event
            merged.append((event, centroid))
        else:
            last_event, last_centroid = merged[-1]
            # Check for overlap
            if (last_event[0] < event[1]) and (last_event[1] > event[0]):
                # Calculate the midpoint of the intervals
                last_event_midpoint = (last_event[0] + last_event[1]) / 2
                current_event_midpoint = (event[0] + event[1]) / 2

                # Choose the event whose centroid is closer to its midpoint
                if abs(last_centroid - last_event_midpoint) <= abs(
                    centroid - current_event_midpoint
                ):
                    continue
                else:
                    merged[-1] = (event, centroid)
            else:
                merged.append((event, centroid))

    final_intervals = [event for event, _ in merged]
    return final_intervals


def evaluator(chunks):
    with torch.no_grad():
        samples = []
        for array, start, end in zip(chunks["audio"], chunks["start"], chunks["end"]):
            samples.append(array["array"][start:end])
        inputs = feature_extractor(
            samples,
            return_tensors="pt",
            sampling_rate=SAMPLING_RATE,
        ).to(device)
        logits = model(**inputs).logits
    y_pred_raw = np.array(logits.cpu())
    y_pred = y_pred_raw.argmax(axis=-1)
    prosodic_units = [
        np.array(frames_to_intervals(i)) + start / SAMPLING_RATE
        for i, start in zip(y_pred, chunks["start"])
    ]
    return {
        "y_pred": y_pred,
        "y_pred_logits": y_pred_raw,
        "prosodic_units": prosodic_units,
    }


audio_duration_samples = (
    Audio(SAMPLING_RATE, mono=True)
    .decode_example({"path": f, "bytes": None})["array"]
    .shape[0]
)
chunk_starts = np.arange(
    0, audio_duration_samples, CHUNK_LENGTH_SAMPLES - OVERLAP_SAMPLES
)
chunk_ends = chunk_starts + CHUNK_LENGTH_SAMPLES

ds = Dataset.from_dict(
    {
        "audio": [f for i in chunk_starts],
        "start": chunk_starts,
        "end": chunk_ends,
        "chunk_centroid_s": (chunk_starts + chunk_ends) / 2 / SAMPLING_RATE,
    }
).cast_column("audio", Audio(SAMPLING_RATE, mono=True))

ds = ds.map(evaluator, batched=True, batch_size=10)


final_intervals = merge_events(ds["prosodic_units"], ds["chunk_centroid_s"])
print(final_intervals)
# Outputs: [[3.14, 4.96], [5.6, 8.4], [8.62, 9.32], [10.12, 10.7], [11.72, 13.1],....
```

## Training Details

| hyperparameter              | value |
| --------------------------- | ----- |
| learning rate               | 3e-5  |
| batch size                  | 1     |
| gradient accumulation steps | 16    |
| num train epochs            | 20    |
| weight decay                | 0.01  |

Software environment can be found in mamba/conda [environment export yml
file](transformers_env.yml). To recreate the environment with conda/mamba, run
`mamba create -f transformers_env.yml` (replace mamba with conda if you don't
use mamba).