import sentencepiece as stp

tokenizer_path = '/tmp/kobert/kobert-news-wiki.spiece'
tokenizer = stp.SentencePieceProcessor(model_file=tokenizer_path)

print("Encode         :", tokenizer.Encode("치킨은 맛있다"))
print("EncodeAsIds    :", tokenizer.EncodeAsIds("치킨은 맛있다"))
print("EncodeAsPieces :", tokenizer.EncodeAsPieces("치킨은 맛있다"))
print("Decode         :", tokenizer.Decode(tokenizer.Encode("치킨")))

import gluonnlp as gnlp

tokenizer_path = '/tmp/kobert/kobert-news-wiki.spiece'
vocab = gnlp.vocab.BERTVocab.from_sentencepiece(
    tokenizer_path, padding_token="[PAD]"
)
tokenizer = gnlp.data.BERTSPTokenizer(tokenizer_path, vocab, lower=False)

tokens = tokenizer("치킨은 맛있다")
ids = tokenizer.convert_tokens_to_ids(tokens)
decodes = vocab.to_tokens(ids)

print("EncodeAsPieces:", tokens)
print("EncodeAsIds   :", ids)
print("Decode        :", decodes)

transform = gnlp.data.BERTSentenceTransform(
    tokenizer,  # gluonnlp.data.transforms.BERTSPTokenizer
    max_seq_length=64,  # 문장의 길이
    pad=True,
    pair=False,
)

text = "하나님이 세상을 이처럼 사랑하사 독생자를 주셨으니 이는 그를 믿는 자마다 멸망하지 않고 영생을 얻게 하려 하심이라"
token_ids, valid_length, segment_ids = transform([text])
print('[token_ids]\n', token_ids)
print('\n[valid_length]\n', valid_length)
print('\n[segment_ids]\n', segment_ids)
print('\n[id -> token]\n', tokenizer.vocab.to_tokens(token_ids.tolist()))

from transformers import ElectraForSequenceClassification
device = torch.device("cpu")
model = ElectraForSequenceClassification.from_pretrained("monologg/koelectra-base-v3-discriminator").to(device)

activation = {}
def get_activation(name):
    def hook(model, input, output):
        activation[name] = output.detach()
    return hook
getattr(model.electra.encoder.layer, '11').output.register_forward_hook(get_activation('output'))