using the training checkpoint for inference

[himaivan]

Hi @lucadellalib,

Thank you for the excellent work.
I have used the training recipe in speechbrain. to train my own model
I obtain multiple models in the ckpt folder:
CKPT.yaml brain.ckpt compressor.ckpt counter.ckpt dataloader-TRAIN.ckpt decompressor.ckpt optimizer.ckpt quantizer.ckpt scheduler.ckpt

How to perform the inference from this model, i.e., encode audio into tokens?
In the example, the codec model is loaded from torch.hub.load using safetensor format. How can I convert the model in the ckpt folder to the safetensor format so it can be loaded with torch.hub.load?

Thank you.

[lucadellalib]

Hi @himaivan ,

The easiest approach is to instantiate the codec using the same modified configuration you trained with, load the individual checkpoints, and then export the model using the built-in method.

For example:

pip install "focalcodec@git+https://github.com/lucadellalib/focalcodec.git@main#egg=focalcodec"

import torch
import focalcodec

original_focalcodec = focalcodec.FocalCodec.from_pretrained("lucadellalib/focalcodec_50hz")

config = "myconfig.json"
codec = focalcodec.FocalCodec.from_config(config)

encoder_checkpoint = original_focalcodec.encoder.state_dict()
compressor_checkpoint = torch.load("path/to/compressor.ckpt", map_location="cpu")
decompressor_checkpoint = torch.load("path/to/decompressor.ckpt", map_location="cpu")
decoder_checkpoint = torch.load("path/to/decoder.ckpt", map_location="cpu")

codec.encoder.load_state_dict(encoder_checkpoint)
codec.compressor.load_state_dict(compressor_checkpoint)
codec.decompressor.load_state_dict(decompressor_checkpoint)
codec.decoder.load_state_dict(decoder_checkpoint)

codec.to_pretrained("mymodel")

For a template of myconfig.json, you can refer to: https://huggingface.co/lucadellalib/focalcodec_50hz/blob/main/config.json

After that, you should be able to load it in the same way as the pretrained FocalCodec checkpoints and use it for inference (focalcodec.FocalCodec.from_pretrained("mymodel")).