|
import torch |
|
from torch import nn, Tensor |
|
from typing import Union, Tuple, List, Iterable, Dict |
|
import logging |
|
import gzip |
|
from tqdm import tqdm |
|
import numpy as np |
|
import os |
|
import json |
|
from ..util import import_from_string, fullname, http_get |
|
from .tokenizer import WordTokenizer, WhitespaceTokenizer |
|
|
|
|
|
class CNN(nn.Module): |
|
"""CNN-layer with multiple kernel-sizes over the word embeddings""" |
|
|
|
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.config_keys = ['in_word_embedding_dimension', 'out_channels', 'kernel_sizes'] |
|
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 tokenize(self, text: str) -> List[int]: |
|
raise NotImplementedError() |
|
|
|
def save(self, output_path: str): |
|
with open(os.path.join(output_path, 'cnn_config.json'), 'w') as fOut: |
|
json.dump(self.get_config_dict(), fOut, indent=2) |
|
|
|
torch.save(self.state_dict(), os.path.join(output_path, 'pytorch_model.bin')) |
|
|
|
def get_config_dict(self): |
|
return {key: self.__dict__[key] for key in self.config_keys} |
|
|
|
@staticmethod |
|
def load(input_path: str): |
|
with open(os.path.join(input_path, 'cnn_config.json'), 'r') as fIn: |
|
config = json.load(fIn) |
|
|
|
weights = torch.load(os.path.join(input_path, 'pytorch_model.bin'), map_location=torch.device('cpu')) |
|
model = CNN(**config) |
|
model.load_state_dict(weights) |
|
return model |
|
|
|
|
