# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3 (ipykernel)
#     language: python
#     name: python3
# ---

# + [markdown] cellId="zl5k218obd8jdw11dabtlo"
# # Свёрточные нейросети и POS-теггинг
#
# POS-теггинг - определение частей речи (снятие частеречной неоднозначности)
# -

# ## Краткие выводы по проведенным экспериментам
#
# ### Лучшие настройки для кастомной модели SentenceLevelPOSTagger:
# Настройки: EPOCH_N = 100, EMB_SIZE = 128, LAYERS_N = 5, KERNELS = [3, 5, 5, 5, 5], DILATIONS = [1, 2, 2, 3, 3], DROPOUT = 0.6, BATCH_NORMALIZATION = False, CRITERION = F.cross_entropy \
# Количество параметров: 775570
#
# ### Оценка качества для лучшего пайплайна:
# - среднее значение функции потерь на обучении: 0.0047452752478420734
# - macro avg f1 на обучении: 0.96
# - среднее значение функции потерь на валидации: 0.010289125144481659
# - macro avg f1 на валидации: 0.93
# - примерное время на одну эпоху (в секундах): 49
#
# Результаты применения теггеров:\
# <i>мама-NOUN мыла-NOUN раму-NOUN</i>
#
# <i>косил-VERB косой-NOUN косой-NOUN косой-NOUN</i>
#
# <i>глокая-ADJ куздра-NOUN штеко-ADV будланула-VERB бокра-NOUN и-CCONJ куздрячит-VERB бокрёнка-NOUN</i>
#
# <i>сяпала-VERB калуша-NOUN с-ADP калушатами-NOUN по-ADP напушке-NOUN</i>
#
# <i>пирожки-NOUN поставлены-VERB в-ADP печь-NOUN мама-NOUN любит-VERB печь-NOUN</i>
#
# <i>ведро-ADV дало-NOUN течь-NOUN вода-NOUN стала-VERB течь-NOUN</i>
#
# <i>три-NUM да-PART три-NUM будет-VERB дырка-NOUN</i>
#
# <i>три-NUM да-PART три-NUM будет-VERB шесть-NUM</i>
#
# <i>сорок-NUM сорок-NUM</i>
#
# ### Резюме:
#
# 1) Увеличение размера эмбеддинга немного повысило качество модели. \
# 2) Для получения высокого качества модели количество эпох следует взять не менее 50. \
# 3) Оптимальное число слоев для реализованной архитектуры StackedConv1d: 5-7. \
# 4) Увеличение силы dropout (0.5-0.7) позволяет избегать переобучения при увеличении количества слоев. \
# 5) Увеличение размера ядра свертки (до 5) на последних слоях немного повысило качество модели. \
# 6) При замене dropout на batch-нормализацию качество модели снизилось. \
# 7) Добавление dilated-сверток повышает качество модели. \
# 8) При использовании взвешивания классов качество модели значительно снижается. \
# 9) Использование отдельного id для обозначения начала и конца токена и обучение для него отдельного эмбеддинга значительно повышают качество модели.

# + cellId="19p4wa36kdus1r68ra7aw"
# Если Вы запускаете ноутбук на colab или kaggle,
# выполните следующие строчки, чтобы подгрузить библиотеку dlnlputils:

# # !git clone https://github.com/arturburiev/stepik-dl-nlp.git
# # !pip install -r stepik-dl-nlp/requirements.txt
# import sys; sys.path.append('./stepik-dl-nlp')

# + cellId="ln9gxth234xc5885po5s"
# # !pip install pyconll
# # !pip install spacy_udpipe

# + cellId="ribvc3xjhkce3pulr6w2c8"
# #!g1.1
# %load_ext autoreload
# %autoreload 2

import sys; sys.path.append('./stepik-dl-nlp')

import warnings
warnings.filterwarnings('ignore')

from sklearn.metrics import classification_report

import numpy as np

import pyconll

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import TensorDataset

import dlnlputils
from dlnlputils.data import tokenize_corpus, build_vocabulary, \
    character_tokenize, pos_corpus_to_tensor, POSTagger
from dlnlputils.pipeline import train_eval_loop, predict_with_model, init_random_seed

init_random_seed()

# + [markdown] cellId="xfuyko841oepayif42tcu"
# ## Загрузка текстов и разбиение на обучающую и тестовую подвыборки

# + cellId="mlowe7zvpf3efhyvxozcn"
# # Если Вы запускаете ноутбук на colab или kaggle, добавьте в начало пути ./stepik-dl-nlp
# # !wget -O ./stepik-dl-nlp/datasets/ru_syntagrus-ud-train.conllu https://raw.githubusercontent.com/UniversalDependencies/UD_Russian-SynTagRus/master/ru_syntagrus-ud-train-a.conllu
# # !wget -O ./stepik-dl-nlp/datasets/ru_syntagrus-ud-dev.conllu https://raw.githubusercontent.com/UniversalDependencies/UD_Russian-SynTagRus/master/ru_syntagrus-ud-dev.conllu

# + cellId="g4a6z2kyqhqwlut62wpqyc"
# #!c1.8
# Если Вы запускаете ноутбук на colab или kaggle, добавьте в начало пути ./stepik-dl-nlp
full_train = pyconll.load_from_file('./stepik-dl-nlp/datasets/ru_syntagrus-ud-train.conllu')
full_test = pyconll.load_from_file('./stepik-dl-nlp/datasets/ru_syntagrus-ud-dev.conllu')

# + cellId="c5dgoa1qgwef4f6nd5khd"
# #!c1.8
for sent in full_train[:2]:
    for token in sent:
        print(token.form, token.upos)
    print()

# + cellId="tlgfduofakbnva1nlh5dse"
# #!c1.8
MAX_SENT_LEN = max(len(sent) for sent in full_train)
MAX_ORIG_TOKEN_LEN = max(len(token.form) for sent in full_train for token in sent)
print('Наибольшая длина предложения', MAX_SENT_LEN)
print('Наибольшая длина токена', MAX_ORIG_TOKEN_LEN)

# + cellId="yhmktzo5wfc0y8xqcuaa66"
# #!c1.8
all_train_texts = [' '.join(token.form for token in sent) for sent in full_train]
print('\n'.join(all_train_texts[:10]))

# + cellId="gdhusetk81nkeq55bq44p"
# #!g1.1
train_char_tokenized = tokenize_corpus(all_train_texts, tokenizer=character_tokenize)
char_vocab, word_doc_freq = build_vocabulary(train_char_tokenized, max_doc_freq=1.0, min_count=5, pad_word='<PAD>')
print("Количество уникальных символов", len(char_vocab))
print(list(char_vocab.items())[:10])

# + cellId="ijh7ya0lvlhibpj111g2k"
# #!g1.1
UNIQUE_TAGS = ['<NOTAG>'] + sorted({token.upos for sent in full_train for token in sent if token.upos})
label2id = {label: i for i, label in enumerate(UNIQUE_TAGS)}
label2id

# + cellId="8sxafylg9liuji00tinwu"
# #!c1.8
START_WORD_ID = len(char_vocab)
END_WORD_ID = len(char_vocab)

# + cellId="uznkra4wkqwow0u2hsy7"
# #!c1.8
train_inputs, train_labels = pos_corpus_to_tensor(full_train, char_vocab, label2id, MAX_SENT_LEN, MAX_ORIG_TOKEN_LEN, start_word_id = START_WORD_ID, end_word_id = END_WORD_ID)
train_dataset = TensorDataset(train_inputs, train_labels)

test_inputs, test_labels = pos_corpus_to_tensor(full_test, char_vocab, label2id, MAX_SENT_LEN, MAX_ORIG_TOKEN_LEN, start_word_id = START_WORD_ID, end_word_id = END_WORD_ID)
test_dataset = TensorDataset(test_inputs, test_labels)

# + cellId="3mtl1tt2i4398aeijhhkci"
# #!c1.8
train_inputs[1][:5]

# + cellId="3wo3f3ujv37hvdsyexkjua"
# #!c1.8
train_labels[1]


# + cellId="65gdv9uuh5qst41zvx8g3" jupyter={"source_hidden": true}
# # #!c1.8
# torch.save(train_inputs, './stepik-dl-nlp/datasets/train_inputs_dataset_postag.pt')
# torch.save(train_labels, './stepik-dl-nlp/datasets/train_labels_dataset_postag.pt')
# torch.save(test_inputs, './stepik-dl-nlp/datasets/test_inputs_dataset_postag.pt')
# torch.save(test_labels, './stepik-dl-nlp/datasets/test_labels_dataset_postag.pt')

# import pickle

# with open('./stepik-dl-nlp/datasets/char_vocab_postag.pt', 'wb') as f:
#     pickle.dump(char_vocab, f)
    
# with open('./stepik-dl-nlp/datasets/label2id_postag.pt', 'wb') as f:
#     pickle.dump(label2id, f)

# + cellId="6a7p15scy7addiwoecv1d" jupyter={"source_hidden": true}
# # #!c1.8
# train_inputs = torch.load('./stepik-dl-nlp/datasets/train_inputs_dataset_postag.pt')
# train_labels = torch.load('./stepik-dl-nlp/datasets/train_labels_dataset_postag.pt')
# test_inputs = torch.load('./stepik-dl-nlp/datasets/test_inputs_dataset_postag.pt')
# test_labels = torch.load('./stepik-dl-nlp/datasets/test_labels_dataset_postag.pt')

# train_dataset = TensorDataset(train_inputs, train_labels)
# test_dataset = TensorDataset(test_inputs, test_labels)

# import pickle

# with open('./stepik-dl-nlp/datasets/char_vocab_postag.pt', 'rb') as f:
#     char_vocab = pickle.load(f)

# # with open('./stepik-dl-nlp/datasets/label2id_postag.pt', 'rb') as file:
# #     label2id = pickle.load(file)

# + [markdown] cellId="f81xb8fuwno6isu29qqal2"
# ## Вспомогательная свёрточная архитектура

# + cellId="jqupmvergbib8tw4ys3rn"
# #!c1.8
class StackedConv1d(nn.Module):
    def __init__(self, features_num, layers_n=1, kernels=3, conv_layer=nn.Conv1d, dropout=0.0, dilations=1):
        super().__init__()
        layers = []
        for i in range(layers_n):
            l_kernel = kernels if isinstance(kernels, int) else kernels[i]
            l_dilation = dilations if isinstance(dilations, int) else dilations[i]
            l_padding = (l_kernel - 1) * l_dilation//2
            
            layers.append(nn.Sequential(
                conv_layer(features_num, features_num, l_kernel, padding=l_padding, dilation=l_dilation),
                nn.Dropout(dropout),
#                 nn.BatchNorm1d(features_num, affine=True),
                nn.LeakyReLU(),
            ))
            
        self.layers = nn.ModuleList(layers)
    
    def forward(self, x):
        """x - BatchSize x FeaturesNum x SequenceLen"""
        for layer in self.layers:
            x = x + layer(x)
        return x


# + [markdown] cellId="2842v34uxaw5oa7advionb" jupyter={"source_hidden": true}
# ## Предсказание частей речи на уровне отдельных токенов

# + cellId="33p33xioysrspskijc2vuk" jupyter={"source_hidden": true}
# #!c1.8
class SingleTokenPOSTagger(nn.Module):
    def __init__(self, vocab_size, labels_num, embedding_size=32, **kwargs):
        super().__init__()
        self.char_embeddings = nn.Embedding(vocab_size, embedding_size, padding_idx=0)
        self.backbone = StackedConv1d(embedding_size, **kwargs)
        self.global_pooling = nn.AdaptiveMaxPool1d(1)
        self.out = nn.Linear(embedding_size, labels_num)
        self.labels_num = labels_num
    
    def forward(self, tokens):
        """tokens - BatchSize x MaxSentenceLen x MaxTokenLen"""
        batch_size, max_sent_len, max_token_len = tokens.shape
        tokens_flat = tokens.view(batch_size * max_sent_len, max_token_len)
        
        char_embeddings = self.char_embeddings(tokens_flat)  # BatchSize*MaxSentenceLen x MaxTokenLen x EmbSize
        char_embeddings = char_embeddings.permute(0, 2, 1)  # BatchSize*MaxSentenceLen x EmbSize x MaxTokenLen
        
        features = self.backbone(char_embeddings)
        
        global_features = self.global_pooling(features).squeeze(-1)  # BatchSize*MaxSentenceLen x EmbSize
        
        logits_flat = self.out(global_features)  # BatchSize*MaxSentenceLen x LabelsNum
        logits = logits_flat.view(batch_size, max_sent_len, self.labels_num)  # BatchSize x MaxSentenceLen x LabelsNum
        logits = logits.permute(0, 2, 1)  # BatchSize x LabelsNum x MaxSentenceLen
        return logits


# + cellId="a7rgombh7hbqgo0s9j6ldk" jupyter={"source_hidden": true}
# #!g1.1
single_token_model = SingleTokenPOSTagger(len(char_vocab), len(label2id), embedding_size=64, layers_n=3, kernel_size=3, dropout=0.3)
print('Количество параметров', sum(np.product(t.shape) for t in single_token_model.parameters()))

# + cellId="7yb1nu92l1hc7yiv4toxov" jupyter={"source_hidden": true}
# #!g1.1
(best_val_loss,
 best_single_token_model) = train_eval_loop(single_token_model,
                                            train_dataset,
                                            test_dataset,
                                            F.cross_entropy,
                                            lr=5e-3,
                                            epoch_n=10,
                                            batch_size=64,
                                            device='cuda',
                                            early_stopping_patience=5,
                                            max_batches_per_epoch_train=500,
                                            max_batches_per_epoch_val=100,
                                            lr_scheduler_ctor=lambda optim: torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=2,
                                                                                                                       factor=0.5,
                                                                                                                       verbose=True))

# + cellId="wf7kngcljcgvqmko9zt0qr" jupyter={"source_hidden": true}
# #!c1.8
# Если Вы запускаете ноутбук на colab или kaggle, добавьте в начало пути ./stepik-dl-nlp
# torch.save(best_single_token_model.state_dict(), './models/single_token_pos.pth')

# + cellId="zhjqmi2cf2agp2q3cv0dww" jupyter={"source_hidden": true}
# #!c1.8
# Если Вы запускаете ноутбук на colab или kaggle, добавьте в начало пути ./stepik-dl-nlp
# single_token_model.load_state_dict(torch.load('./models/single_token_pos.pth'))

# + cellId="32dzhka5lrqpje57n1aj9" jupyter={"source_hidden": true}
# #!g1.1
single_token_model = best_single_token_model

# + cellId="gc2olr6cavrot8t0qx634" jupyter={"source_hidden": true}
# #!g1.1
train_pred = predict_with_model(single_token_model, train_dataset)
train_loss = F.cross_entropy(torch.tensor(train_pred),
                             torch.tensor(train_labels))
print('Среднее значение функции потерь на обучении', float(train_loss))
print(classification_report(train_labels.view(-1), train_pred.argmax(1).reshape(-1), target_names=UNIQUE_TAGS))
print()

test_pred = predict_with_model(single_token_model, test_dataset)
test_loss = F.cross_entropy(torch.tensor(test_pred),
                            torch.tensor(test_labels))
print('Среднее значение функции потерь на валидации', float(test_loss))
print(classification_report(test_labels.view(-1), test_pred.argmax(1).reshape(-1), target_names=UNIQUE_TAGS))


# + [markdown] cellId="s0yc7vqix3lnaqbd0dx4i"
# ## Предсказание частей речи на уровне предложений (с учётом контекста)

# + cellId="gns1qaer0pg1kkkub83w"
# #!c1.8
class SentenceLevelPOSTagger(nn.Module):
    def __init__(self, vocab_size, labels_num, embedding_size=32, single_backbone_kwargs={}, context_backbone_kwargs={}):
        super().__init__()
        self.embedding_size = embedding_size
        self.char_embeddings = nn.Embedding(vocab_size, embedding_size, padding_idx=0)
        self.single_token_backbone = StackedConv1d(embedding_size, **single_backbone_kwargs)
        self.context_backbone = StackedConv1d(embedding_size, **context_backbone_kwargs)
        self.global_pooling = nn.AdaptiveMaxPool1d(1)
        self.out = nn.Conv1d(embedding_size, labels_num, 1)
        self.labels_num = labels_num
    
    def forward(self, tokens):
        """tokens - BatchSize x MaxSentenceLen x MaxTokenLen"""
        batch_size, max_sent_len, max_token_len = tokens.shape
        tokens_flat = tokens.view(batch_size * max_sent_len, max_token_len)
        
        char_embeddings = self.char_embeddings(tokens_flat)  # BatchSize*MaxSentenceLen x MaxTokenLen x EmbSize
        char_embeddings = char_embeddings.permute(0, 2, 1)  # BatchSize*MaxSentenceLen x EmbSize x MaxTokenLen
        char_features = self.single_token_backbone(char_embeddings)
        
        token_features_flat = self.global_pooling(char_features).squeeze(-1)  # BatchSize*MaxSentenceLen x EmbSize

        token_features = token_features_flat.view(batch_size, max_sent_len, self.embedding_size)  # BatchSize x MaxSentenceLen x EmbSize
        token_features = token_features.permute(0, 2, 1)  # BatchSize x EmbSize x MaxSentenceLen
        context_features = self.context_backbone(token_features)  # BatchSize x EmbSize x MaxSentenceLen

        logits = self.out(context_features)  # BatchSize x LabelsNum x MaxSentenceLen
        return logits


# + cellId="gskyshgnachfvqly54448q"
# #!g1.1

# Изначальные настройки
# EPOCH_N = 10
# EMB_SIZE = 64
# DROPOUT = 0.3
# LAYERS_N = 3
# KERNELS = 3
# DILATIONS = 1
# BATCH_NORMALIZATION = False

# Итоговые настройки
EPOCH_N = 100
EMB_SIZE = 128
DROPOUT = 0.6
LAYERS_N = 5
KERNELS = [3, 5, 5, 5, 5]
DILATIONS = [1, 2, 2, 3, 3]
BATCH_NORMALIZATION = False

sentence_level_model = SentenceLevelPOSTagger(len(char_vocab) + 1, len(label2id), embedding_size=EMB_SIZE,
                                              single_backbone_kwargs=dict(layers_n=LAYERS_N, kernels=KERNELS, dropout=DROPOUT, dilations=DILATIONS),
                                              context_backbone_kwargs=dict(layers_n=LAYERS_N, kernels=KERNELS, dropout=DROPOUT, dilations=DILATIONS))

print(f"Настройки: EPOCH_N = {EPOCH_N}, EMB_SIZE = {EMB_SIZE}, LAYERS_N = {LAYERS_N}, KERNELS = {KERNELS}, DILATIONS = {DILATIONS}, DROPOUT = {DROPOUT}, BATCH_NORMALIZATION = {BATCH_NORMALIZATION}")
print('Количество параметров:', sum(np.product(t.shape) for t in sentence_level_model.parameters()))

# + [markdown] cellId="7n21omaha3np7j2zfdx46a" jupyter={"source_hidden": true}
# ### Взвешивание классов

# + cellId="efd6t04b1t740cv775fi2t" jupyter={"source_hidden": true}
# #!c1.8
from sklearn.utils.class_weight import compute_class_weight
class_weights=compute_class_weight('balanced',np.unique(train_labels),train_labels.numpy().ravel())
class_weights=torch.tensor(class_weights,dtype=torch.float)

# + cellId="tb0hife7osg0bas06fgrj"
# #!g1.1
import sys; sys.path.append('./stepik-dl-nlp')
from dlnlputils.pipeline import train_eval_loop, predict_with_model, init_random_seed

torch.backends.cudnn.benchmark=True
torch.backends.cudnn.deterministic=False 

criterion_func = F.cross_entropy
# criterion_func = nn.CrossEntropyLoss(weight=class_weights,reduction='mean').cuda()

(best_val_loss,
 best_sentence_level_model) = train_eval_loop(sentence_level_model,
                                              train_dataset,
                                              test_dataset,
                                              criterion_func,
                                              lr=5e-3,
                                              epoch_n=EPOCH_N,
                                              batch_size=64,
                                              device='cuda',
                                              early_stopping_patience=5,
                                              max_batches_per_epoch_train=500,
                                              max_batches_per_epoch_val=100,
                                              lr_scheduler_ctor=lambda optim: torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=2,
                                                                                                                         factor=0.5,
                                                                                                                         verbose=True))

# + cellId="5r98nrn1mfp6w0l7qzrgou" jupyter={"source_hidden": true}
# Если Вы запускаете ноутбук на colab или kaggle, добавьте в начало пути ./stepik-dl-nlp
# torch.save(best_sentence_level_model.state_dict(), './models/sentence_level_pos.pth')

# + cellId="p9vaps0ic1l86x3u4oiggd" jupyter={"source_hidden": true}
# Если Вы запускаете ноутбук на colab или kaggle, добавьте в начало пути ./stepik-dl-nlp
# sentence_level_model.load_state_dict(torch.load('./models/sentence_level_pos.pth'))

# + cellId="8e075fh07hvmhjy6qas2do"
# #!g1.1
sentence_level_model = best_sentence_level_model

# + cellId="nqrx2pwd46itu94x7jshz"
# #!g1.1
UNIQUE_TAGS = list(label2id.keys())
train_pred = predict_with_model(sentence_level_model, train_dataset)
train_loss = F.cross_entropy(torch.tensor(train_pred),
                             train_labels.clone().detach())
print('Среднее значение функции потерь на обучении', float(train_loss))
print(classification_report(train_labels.view(-1), train_pred.argmax(1).reshape(-1), target_names=UNIQUE_TAGS))
print()

test_pred = predict_with_model(sentence_level_model, test_dataset)
test_loss = F.cross_entropy(torch.tensor(test_pred),
                            test_labels.clone().detach())
print('Среднее значение функции потерь на валидации', float(test_loss))
print(classification_report(test_labels.view(-1), test_pred.argmax(1).reshape(-1), target_names=UNIQUE_TAGS))

# + [markdown] cellId="uuw6k0pj8wevchahu90ugg"
# ## Применение полученных теггеров и сравнение

# +
# #!g1.1
single_token_pos_tagger = POSTagger(sentence_level_model, char_vocab, UNIQUE_TAGS, MAX_SENT_LEN, MAX_ORIG_TOKEN_LEN, start_word_id = START_WORD_ID, end_word_id = END_WORD_ID)

test_sentences = [
    '<NAME>.',
    'Косил косой косой косой.',
    'Глокая куздра штеко будланула бокра и куздрячит бокрёнка.',
    'Сяпала Калуша с Калушатами по напушке.',
    'Пирожки поставлены в печь, мама любит печь.',
    'Ведро дало течь, вода стала течь.',
    'Три да три, будет дырка.',
    'Три да три, будет шесть.',
    'Сорок сорок'
]
test_sentences_tokenized = tokenize_corpus(test_sentences, min_token_size=1)

for sent_tokens, sent_tags in zip(test_sentences_tokenized, single_token_pos_tagger(test_sentences)):
    print(' '.join('{}-{}'.format(tok, tag) for tok, tag in zip(sent_tokens, sent_tags)))
    print()

# + cellId="fntk3o3pt5m569uyiqdgdk"
# #!g1.1
sentence_level_pos_tagger = POSTagger(sentence_level_model, char_vocab, UNIQUE_TAGS, MAX_SENT_LEN, MAX_ORIG_TOKEN_LEN, start_word_id = START_WORD_ID, end_word_id = END_WORD_ID)

test_sentences = [
    '<NAME>.',
    'Косил косой косой косой.',
    'Глокая куздра штеко будланула бокра и куздрячит бокрёнка.',
    'Сяпала Калуша с Калушатами по напушке.',
    'Пирожки поставлены в печь, мама любит печь.',
    'Ведро дало течь, вода стала течь.',
    'Три да три, будет дырка.',
    'Три да три, будет шесть.',
    'Сорок сорок'
]
test_sentences_tokenized = tokenize_corpus(test_sentences, min_token_size=1)

for sent_tokens, sent_tags in zip(test_sentences_tokenized, sentence_level_pos_tagger(test_sentences)):
    print(' '.join('{}-{}'.format(tok, tag) for tok, tag in zip(sent_tokens, sent_tags)))
    print()


# + [markdown] cellId="ufj5u78t79x8g7s5iigyl" jupyter={"source_hidden": true}
# ## Свёрточный модуль своими руками

# + cellId="zzgwdt9xebmo1etxwus0uh" jupyter={"source_hidden": true}
class MyConv1d(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, padding=0):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.padding = padding
        self.weight = nn.Parameter(torch.randn(in_channels * kernel_size, out_channels) / (in_channels * kernel_size),
                                   requires_grad=True)
        self.bias = nn.Parameter(torch.zeros(out_channels), requires_grad=True)
    
    def forward(self, x):
        """x - BatchSize x InChannels x SequenceLen"""

        batch_size, src_channels, sequence_len = x.shape        
        if self.padding > 0:
            pad = x.new_zeros(batch_size, src_channels, self.padding)
            x = torch.cat((pad, x, pad), dim=-1)
            sequence_len = x.shape[-1]

        chunks = []
        chunk_size = sequence_len - self.kernel_size + 1
        for offset in range(self.kernel_size):
            chunks.append(x[:, :, offset:offset + chunk_size])

        in_features = torch.cat(chunks, dim=1)  # BatchSize x InChannels * KernelSize x ChunkSize
        in_features = in_features.permute(0, 2, 1)  # BatchSize x ChunkSize x InChannels * KernelSize
        out_features = torch.bmm(in_features, self.weight.unsqueeze(0).expand(batch_size, -1, -1)) + self.bias.unsqueeze(0).unsqueeze(0)
        out_features = out_features.permute(0, 2, 1)  # BatchSize x OutChannels x ChunkSize
        return out_features


# + cellId="xn38fh3coni823aqd8bkj6" jupyter={"source_hidden": true}
sentence_level_model_my_conv = SentenceLevelPOSTagger(len(char_vocab), len(label2id), embedding_size=64,
                                                      single_backbone_kwargs=dict(layers_n=3, kernel_size=3, dropout=0.3, conv_layer=MyConv1d),
                                                      context_backbone_kwargs=dict(layers_n=3, kernel_size=3, dropout=0.3, conv_layer=MyConv1d))
print('Количество параметров', sum(np.product(t.shape) for t in sentence_level_model_my_conv.parameters()))

# + cellId="86q7sp6ipmg8t8rklgvw3g" jupyter={"source_hidden": true}
(best_val_loss,
 best_sentence_level_model_my_conv) = train_eval_loop(sentence_level_model_my_conv,
                                                      train_dataset,
                                                      test_dataset,
                                                      F.cross_entropy,
                                                      lr=5e-3,
                                                      epoch_n=10,
                                                      batch_size=64,
                                                      device='cuda',
                                                      early_stopping_patience=5,
                                                      max_batches_per_epoch_train=500,
                                                      max_batches_per_epoch_val=100,
                                                      lr_scheduler_ctor=lambda optim: torch.optim.lr_scheduler.ReduceLROnPlateau(optim, patience=2,
                                                                                                                                 factor=0.5,
                                                                                                                                 verbose=True))

# + cellId="4ac0m56m8x9kq8bbs9v0u" jupyter={"source_hidden": true}
train_pred = predict_with_model(best_sentence_level_model_my_conv, train_dataset)
train_loss = F.cross_entropy(torch.tensor(train_pred),
                             torch.tensor(train_labels))
print('Среднее значение функции потерь на обучении', float(train_loss))
print(classification_report(train_labels.view(-1), train_pred.argmax(1).reshape(-1), target_names=UNIQUE_TAGS))
print()

test_pred = predict_with_model(best_sentence_level_model_my_conv, test_dataset)
test_loss = F.cross_entropy(torch.tensor(test_pred),
                            torch.tensor(test_labels))
print('Среднее значение функции потерь на валидации', float(test_loss))
print(classification_report(test_labels.view(-1), test_pred.argmax(1).reshape(-1), target_names=UNIQUE_TAGS))
# -



StarcoderdataJupyterScriptsDedupFiltered