huggingface / diarizers

246 stars 16 forks source link

How to save the finetuned model as a .bin file? #4

Closed utility-aagrawal closed 5 months ago

utility-aagrawal commented 5 months ago

Hi,

I finetuned the pyannote-segmentation model for my usecase but it is saved as a model.safetensors file. Can I convert it to a pytorch_model.bin file? I am using whisperx to create speaker-aware transcripts and .safetensors isn't working with that library. Thanks!

kamilakesbi commented 5 months ago

Hi @utility-aagrawal,

To save the fine-tuned model as a pytorch_model.bin file, you can simply add the argument save_safetensors=False to the TrainingArguments of the Trainer ;)

If you use the train_segmentation.py script of the library, you could run:

python3 train_segmentation.py \
    --dataset_name=diarizers-community/callhome \
    --dataset_config_name=jpn \
    --split_on_subset=data \
    --model_name_or_path=pyannote/segmentation-3.0 \
    --output_dir=./speaker-segmentation-fine-tuned-callhome-jpn \
    --do_train \
    --do_eval \
    --learning_rate=1e-3 \
    --num_train_epochs=5 \
    --lr_scheduler_type=cosine \
    --per_device_train_batch_size=32 \
    --per_device_eval_batch_size=32 \
    --evaluation_strategy=epoch \
    --save_strategy=epoch \
    --preprocessing_num_workers=2 \
    --dataloader_num_workers=2 \
    --logging_steps=100 \
    --load_best_model_at_end \
    --push_to_hub \
    --save_safetensors=False
utility-aagrawal commented 5 months ago

Thanks @kamilakesbi ! I did that and was able to save the model as a pytorch_model.bin file. There's still something different though because I can't get this finetuned model to work with whisperx. My guess is because whisperx is expecting the model in pyannote format. Please see here how whisperx loads the model: https://github.com/m-bain/whisperX/blob/f2da2f858e99e4211fe4f64b5f2938b007827e17/whisperx/diarize.py#L19C9-L19C100

When I look at config.yaml files for the finetuned model vs pyannote/segmentation-3.0 model, they are completely different. Finetuned model's config.yaml: image

pyannote/segmentation-3.0's config.yaml: image

Do you know how to push the model to huggingface hub in pyannote format ? Thanks!

kamilakesbi commented 5 months ago

Hi @utility-aagrawal,

Thanks for the question! the fine-tuned segmentation model (from the SegmentationModel class) comes in a different format than the pyannote one in order to be compatible with the HF Trainer.

You can convert it to a pyannote compatible format by using the to_pyannote_model method:

model = SegmentationModel().from_pretrained("diarizers-community/speaker-segmentation-fine-tuned-callhome-jpn")
model = model.to_pyannote_model()

You can push the SegmentationModel to the Hub, but not the pyannote format directly yet! We plan to add this functionality in the near future.