You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
89 lines
2.8 KiB
89 lines
2.8 KiB
from __future__ import annotations
|
|
|
|
try:
|
|
from typing import Self
|
|
except ImportError:
|
|
from typing_extensions import Self
|
|
|
|
import torch
|
|
from torch import nn
|
|
|
|
from sentence_transformers.models.Module import Module
|
|
|
|
|
|
class CNN(Module):
|
|
"""CNN-layer with multiple kernel-sizes over the word embeddings"""
|
|
|
|
config_keys: list[str] = ["in_word_embedding_dimension", "out_channels", "kernel_sizes"]
|
|
config_file_name: str = "cnn_config.json"
|
|
|
|
def __init__(
|
|
self,
|
|
in_word_embedding_dimension: int,
|
|
out_channels: int = 256,
|
|
kernel_sizes: list[int] = [1, 3, 5],
|
|
stride_sizes: list[int] = None,
|
|
):
|
|
nn.Module.__init__(self)
|
|
self.in_word_embedding_dimension = in_word_embedding_dimension
|
|
self.out_channels = out_channels
|
|
self.kernel_sizes = kernel_sizes
|
|
|
|
self.embeddings_dimension = out_channels * len(kernel_sizes)
|
|
self.convs = nn.ModuleList()
|
|
|
|
in_channels = in_word_embedding_dimension
|
|
if stride_sizes is None:
|
|
stride_sizes = [1] * len(kernel_sizes)
|
|
|
|
for kernel_size, stride in zip(kernel_sizes, stride_sizes):
|
|
padding_size = int((kernel_size - 1) / 2)
|
|
conv = nn.Conv1d(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
kernel_size=kernel_size,
|
|
stride=stride,
|
|
padding=padding_size,
|
|
)
|
|
self.convs.append(conv)
|
|
|
|
def forward(self, features):
|
|
token_embeddings = features["token_embeddings"]
|
|
|
|
token_embeddings = token_embeddings.transpose(1, -1)
|
|
vectors = [conv(token_embeddings) for conv in self.convs]
|
|
out = torch.cat(vectors, 1).transpose(1, -1)
|
|
|
|
features.update({"token_embeddings": out})
|
|
return features
|
|
|
|
def get_word_embedding_dimension(self) -> int:
|
|
return self.embeddings_dimension
|
|
|
|
def save(self, output_path: str, *args, safe_serialization: bool = True, **kwargs) -> None:
|
|
self.save_config(output_path)
|
|
self.save_torch_weights(output_path, safe_serialization=safe_serialization)
|
|
|
|
@classmethod
|
|
def load(
|
|
cls,
|
|
model_name_or_path: str,
|
|
subfolder: str = "",
|
|
token: bool | str | None = None,
|
|
cache_folder: str | None = None,
|
|
revision: str | None = None,
|
|
local_files_only: bool = False,
|
|
**kwargs,
|
|
) -> Self:
|
|
hub_kwargs = {
|
|
"subfolder": subfolder,
|
|
"token": token,
|
|
"cache_folder": cache_folder,
|
|
"revision": revision,
|
|
"local_files_only": local_files_only,
|
|
}
|
|
config = cls.load_config(model_name_or_path=model_name_or_path, **hub_kwargs)
|
|
model = cls(**config)
|
|
model = cls.load_torch_weights(model_name_or_path=model_name_or_path, model=model, **hub_kwargs)
|
|
return model
|