74 lines
3.1 KiB
Python
74 lines
3.1 KiB
Python
|
|
from torch import Tensor, nn
|
|
|
|
class NewTokenEmb(nn.Module):
|
|
"""
|
|
For adding new tokens to a pretrained model
|
|
"""
|
|
|
|
def __init__(self,
|
|
old_embeddings: nn.Embedding,
|
|
new_num_tokens: int = None) -> None:
|
|
|
|
super().__init__()
|
|
|
|
self.num_tokens = old_embeddings.num_embeddings + new_num_tokens
|
|
self.old_num_tokens = old_embeddings.num_embeddings
|
|
self.new_num_tokens = new_num_tokens
|
|
self.embedding_dim = old_embeddings.embedding_dim
|
|
|
|
# For text embeddings
|
|
self.text_embeddings = nn.Embedding(
|
|
self.num_tokens,
|
|
self.embedding_dim,
|
|
device=old_embeddings.weight.device,
|
|
dtype=old_embeddings.weight.dtype)
|
|
with torch.no_grad():
|
|
self.text_embeddings.weight.data[:old_embeddings.
|
|
num_embeddings] = old_embeddings.weight.data
|
|
self.text_embeddings.weight.data[
|
|
self.old_num_tokens:] = torch.zeros(
|
|
self.new_num_tokens,
|
|
self.embedding_dim,
|
|
dtype=old_embeddings.weight.dtype,
|
|
device=old_embeddings.weight.device)
|
|
self.text_embeddings.weight.requires_grad_(False)
|
|
|
|
# For motion embeddings
|
|
self.motion_embeddings = nn.Embedding(
|
|
new_num_tokens,
|
|
self.embedding_dim,
|
|
device=old_embeddings.weight.device,
|
|
dtype=old_embeddings.weight.dtype)
|
|
with torch.no_grad():
|
|
self.motion_embeddings.weight.data[:self.
|
|
old_num_tokens] = torch.zeros(
|
|
new_num_tokens,
|
|
self.embedding_dim,
|
|
dtype=old_embeddings.weight.
|
|
dtype,
|
|
device=old_embeddings.
|
|
weight.device)
|
|
self.word2motionProj = nn.Linear(self.old_num_tokens, new_num_tokens)
|
|
|
|
def forward(self, input: Tensor) -> Tensor:
|
|
|
|
with torch.no_grad():
|
|
self.motion_embeddings.weight.data[:self.
|
|
old_num_tokens] = torch.zeros(
|
|
self.new_num_tokens,
|
|
self.embedding_dim,
|
|
dtype=self.motion_embeddings
|
|
.weight.dtype,
|
|
device=self.
|
|
motion_embeddings.weight.
|
|
device)
|
|
|
|
self.motion_embeddings.weight.data[
|
|
self.old_num_tokens:] = self.word2motionProj(
|
|
self.text_embeddings.weight.data[:self.old_num_tokens].permute(
|
|
1, 0)).permute(1, 0)
|
|
|
|
return self.text_embeddings(input) + self.motion_embeddings(input)
|
|
|