error occurs when using my own trained LASER model

[herowolf]

Hi,
I follow the instructions from LASER https://github.com/facebookresearch/fairseq/blob/nllb/examples/laser/README.md and train a LASER model with my own data (Chinese-English Bitext).
But when I embed sentences with the trained model, an error occurs as follows:

from the logs and codes, I found the checkpoint file of my trained models is different from laser2.pt, and the state_dict of my checkpoint file has no "params".

So is there something wrong when I trained the LASER model (actually I want to train a Teacher model)?
Thanks for your reply.

[Kevin Heffernan]

Hi @weitaizhang, did you train this model using --arch laser_lstm?

[herowolf]

Hi @weitaizhang, did you train this model using --arch laser_lstm?

@heffernankevin yes, I used --arch laser_lstm.
FYI, here's my training script:

And did you train the laser2 model with fairseq codes? Or should I convert the trained checkpoint file to some certain structure after training ?

[Kevin Heffernan]

Hi @weitaizhang, yes you'll need to convert the checkpoint file. Can you try running the following temporary solution script (from the $LASER/source directory) and then try re-running the embed step?

import torch
from fairseq.data.dictionary import Dictionary
from embed import LaserLstmEncoder

ckpt = torch.load([path/to/your/checkpoint], map_location="cpu")

encoder_state = {
    key.replace("encoder.", ""): val
    for key, val in ckpt["model"].items()
    if key.startswith("encoder.")
}

params = {
    "num_embeddings": encoder_state["embed_tokens.weight"].shape[0],
    "padding_idx": 1,
    "embed_dim": encoder_state["embed_tokens.weight"].shape[1],
    "hidden_size": ckpt["args"].encoder_hidden_size,
    "num_layers": ckpt["args"].encoder_layers,
    "bidirectional": ckpt["args"].encoder_bidirectional,
}
encoder = LaserLstmEncoder(**params)
encoder.embed_tokens.weight = torch.nn.Parameter(encoder_state["embed_tokens.weight"])
encoder.lstm.load_state_dict(
    {
        key.replace("lstm.", ""): val
        for key, val in encoder_state.items()
        if key.startswith("lstm.")
    }
)

state_dict = {
    "params": params,
    "model": encoder.state_dict(),
    "dictionary": Dictionary.load([path/to/your/src_vocab/file/from/config.json]).indices,
}
torch.save(state_dict, "encoder.pt")

[herowolf]

Hi @weitaizhang, yes you'll need to convert the checkpoint file. Can you try running the following temporary solution script (from the $LASER/source directory) and then try re-running the embed step?

import torch
from fairseq.data.dictionary import Dictionary
from embed import LaserLstmEncoder

ckpt = torch.load([path/to/your/checkpoint], map_location="cpu")

encoder_state = {
    key.replace("encoder.", ""): val
    for key, val in ckpt["model"].items()
    if key.startswith("encoder.")
}

params = {
    "num_embeddings": encoder_state["embed_tokens.weight"].shape[0],
    "padding_idx": 1,
    "embed_dim": encoder_state["embed_tokens.weight"].shape[1],
    "hidden_size": ckpt["args"].encoder_hidden_size,
    "num_layers": ckpt["args"].encoder_layers,
    "bidirectional": ckpt["args"].encoder_bidirectional,
}
encoder = LaserLstmEncoder(**params)
encoder.embed_tokens.weight = torch.nn.Parameter(encoder_state["embed_tokens.weight"])
encoder.lstm.load_state_dict(
    {
        key.replace("lstm.", ""): val
        for key, val in encoder_state.items()
        if key.startswith("lstm.")
    }
)

state_dict = {
    "params": params,
    "model": encoder.state_dict(),
    "dictionary": Dictionary.load([path/to/your/src_vocab/file/from/config.json]).indices,
}
torch.save(state_dict, "encoder.pt")

Many thanks, it works now！