Getting encoder embeddings for words from the model

Mikey · November 27, 2023, 11:41am

Hi!
I’m not the best coder, so I apologise in advance for any potential mistakes in understanding the architecture or terminology. I have a built .pt translation model and I need somehow to extract from it embeddings for particular words (from the first layer of the encoder, if I’m not mistaken.) Could you please guide me how I can do that using openmt py source code?

Thank you in advance

vince62s · November 27, 2023, 3:57pm

if your point is just to extract the embeddings out of the model and use them outside there is an old script here:

github.com

OpenNMT/OpenNMT-py/blob/master/tools/extract_embeddings.py

import argparse

import torch
import onmt
import onmt.model_builder
from onmt.models.model_saver import load_checkpoint
from onmt.utils.parse import ArgumentParser
import onmt.opts
from onmt.inputters.inputter import dict_to_vocabs

# from onmt.utils.misc import use_gpu
from onmt.utils.logging import init_logger, logger

parser = argparse.ArgumentParser(description="translate.py")

parser.add_argument("-model", required=True, help="Path to model .pt file")
parser.add_argument(
    "-output_dir", default=".", help="""Path to output the embeddings"""
)
parser.add_argument("-gpu", type=int, default=-1, help="Device to run on")

This file has been truncated. show original

if you want to use them within your code to do something else, then identify in the script above how you access those and it should be clearer.

Mikey · November 27, 2023, 5:20pm

My initial plan was to replace the original encoder’s embeddings of specific words with alternative vectors like an average between two embeddings or x2 etc…And to translate with the model using these changed embeddings and to see the results. Is there a way to do that?

vince62s · November 27, 2023, 5:42pm

still not sure what you are trying to do but on the encoder side it is happening here:

github.com

OpenNMT/OpenNMT-py/blob/master/onmt/encoders/transformer.py#L246


      
                  norm_eps=opt.norm_eps,
                  use_ckpting=opt.use_ckpting,
                  parallel_gpu=opt.world_size
                  if opt.parallel_mode == "tensor_parallel"
                  else 1,
                  rotary_interleave=opt.rotary_interleave,
              )
          
          def forward(self, src, src_len=None):
              """See :func:`EncoderBase.forward()`"""
              enc_out = self.embeddings(src)
              mask = sequence_mask(src_len).unsqueeze(1).unsqueeze(1)
              mask = mask.expand(-1, -1, mask.size(3), -1)
              # Padding mask is now (batch x 1 x slen x slen)
              # 1 to be expanded to number of heads in MHA
              # Run the forward pass of every layer of the tranformer.
          
              for layer in self.transformer:
                  enc_out = layer(enc_out, mask)
              enc_out = self.layer_norm(enc_out)
              return enc_out, None, src_len

Mikey · November 30, 2023, 12:09pm

Hi! I am getting this error: ImportError: cannot import name ‘dict_to_vocabs’ from ‘onmt.inputters.inputter’. when I try to run this script. What can be a reason? Thank you in advance!