from pyannote.database import get_protocol
ami = get_protocol('AMI.SpeakerDiarization.only_words')

"""
Explanation: Training, fine-tuning, and transfer learning with pyannote.audio
In this tutorial, you will learn how to use pyannote.audio to 

train a voice activity detection model from scratch, 
fine-tune a pretrained speaker segmentation model,
perform transfer learning (from speaker segmentation to overlapped speech detection)

Data preparation
This tutorial assumes that 
* the AMI corpus has already been setup for use with pyannote
* the PYANNOTE_DATABASE_CONFIG environment variable is set accordingly. 
See pyannote.database documentation to learn how to prepare your own dataset for training with pyannote.audio.
End of explanation
"""


from pyannote.audio.tasks import VoiceActivityDetection
vad_task = VoiceActivityDetection(ami, duration=2.0, batch_size=128)

"""
Explanation: Training a voice activity detection model from scratch
Voice activity detection (VAD) is the task of detecting speech regions in a given audio stream or recording.
We initialize a VAD task that describes how the model will be trained:

ami indicates that we will use files available in ami.train().
duration=2. and batch_size=128 indicates that the model will ingest batches of 128 two seconds long audio chunks.
End of explanation
"""


from pyannote.audio.models.segmentation import PyanNet
vad_model = PyanNet(task=vad_task, sincnet={'stride': 10})

"""
Explanation: We initialize one model with the PyanNet architecture used in that paper.
In particular, we increase the default stride of the initial sincnet feature extraction layer to 10.
The model is also provided with the task (task=vad_task) for which it is being trained:
End of explanation
"""


import pytorch_lightning as pl
trainer = pl.Trainer(gpus=1, max_epochs=1)
trainer.fit(vad_model)

"""
Explanation: Now that everything is ready, let's train with pytorch-ligthning!
End of explanation
"""


test_file = next(ami.test())
# here we use a test file provided by the protocol, but it could be any audio file
# e.g. test_file = "/path/to/test.wav".

"""
Explanation: For the purpose of this tutorial, the model is trained for only 1 epoch. One can obviously expect better performance by training longer and on more data.
See pytorch-lightning documentation to learn more about the Trainer API, in particular.
Once trained, the model can be applied to a test file:
End of explanation
"""


from pyannote.audio import Inference
vad = Inference(vad_model)

vad_probability = vad(test_file)
vad_probability

"""
Explanation: Because the model was trained on 2s audio chunks and that test files are likely to be much longer than that, we wrap the model with an Inference instance: it will take care of sliding a 2s window over the whole file and aggregate the output of the model.
End of explanation
"""


expected_output = test_file["annotation"].get_timeline().support()
expected_output

"""
Explanation: Perfect voice activity detection output should look like that:
End of explanation
"""


from pyannote.audio import Model
pretrained = Model.from_pretrained("pyannote/segmentation")

"""
Explanation: Fine-tuning a pretrained speaker segmentation model
Speaker diarization is the task of partitioning a given audio stream of recording into according to the speaker identity.
pyannote/segmentation is a model that was pretrained to perform speaker diarization, but only locally, on 5s-long audio chunks.
End of explanation
"""


spk_probability = Inference(pretrained, step=2.5)(test_file)
spk_probability

"""
Explanation: Let's visualize how it performs on our test file:
End of explanation
"""


test_file["annotation"].discretize(notebook.crop, resolution=0.010)

"""
Explanation: A perfect output would look like that:
End of explanation
"""


from pyannote.audio.tasks import Segmentation
seg_task = Segmentation(ami, duration=5.0, max_num_speakers=4)

"""
Explanation: We are going to fine-tune this pretrained model on the AMI dataset:
End of explanation
"""


def test(model, protocol, subset="test"):
    from pyannote.audio.utils.signal import binarize
    from pyannote.audio.utils.metric import DiscreteDiarizationErrorRate
    from pyannote.audio.pipelines.utils import get_devices

    (device,) = get_devices(needs=1)
    metric = DiscreteDiarizationErrorRate()
    files = list(getattr(protocol, subset)())

    inference = Inference(model, device=device)

    for file in files:
        reference = file["annotation"]
        hypothesis = binarize(inference(file))
        uem = file["annotated"]
        _ = metric(reference, hypothesis, uem=uem)
        
    return abs(metric)

"""
Explanation: To check that fine-tuning was actually helpful, we need to evaluate the performance of the pretrained model, and compute the average local diarization error rate on a 5s window sliding over the whole test set. To do so, we need to create a helper function:
End of explanation
"""


der_pretrained = test(model=pretrained, protocol=ami, subset="test")
print(f"Local DER (pretrained) = {der_pretrained * 100:.1f}%")

"""
Explanation: We can then evaluate the model and see its local DER:
End of explanation
"""


from copy import deepcopy
finetuned = deepcopy(pretrained)
finetuned.task = seg_task

"""
Explanation: Next, we prepare the model for fine-tuning, simply by overriding its task attribute...
End of explanation
"""


trainer = pl.Trainer(gpus=1, max_epochs=1)
trainer.fit(finetuned)

"""
Explanation: ... and we train it (for just one epoch)
End of explanation
"""


der_finetuned = test(model=finetuned, protocol=ami, subset="test")
print(f"Local DER (finetuned) = {der_finetuned * 100:.1f}%")

"""
Explanation: We now evaluate the performance of the fine-tuned model...
End of explanation
"""


Inference('pyannote/segmentation', step=2.5)(test_file)

test_file["annotation"]

"""
Explanation: ... to confirm that it actually improved its performance on the AMI test set.
Transfer learning
What if you are only interested in detecting overlapped speech regions? 
Looking at the output of the pyannote/segmentation model, it seems that it would be a good starting point for training such a dedicated model:
End of explanation
"""


from pyannote.audio.tasks import OverlappedSpeechDetection
osd_task = OverlappedSpeechDetection(ami, duration=2.0)

osd_model = Model.from_pretrained("pyannote/segmentation")
osd_model.task = osd_task

"""
Explanation: Let's update the model so that it specifically addresses overlapped speech detection.
This is achieved very simply by updating the task attribute...
End of explanation
"""


osd_model.freeze_up_to('lstm')

"""
Explanation: ... optionally freeezing a bunch of layers...
End of explanation
"""


trainer = pl.Trainer(gpus=1, max_epochs=1)
trainer.fit(osd_model)

"""
Explanation: ... and training it:
End of explanation
"""


from pyannote.audio.utils.signal import binarize
binarize(Inference(osd_model)(test_file))

test_file["annotation"]

"""
Explanation: Et voilà! A brand new overlapped speech detection model!
End of explanation
"""


