from ditk import logging

from gende.train import train

if __name__ == '__main__':
    logging.try_init_root(logging.INFO)
    train(
        'runs/bert_mean_seed_0',
        model_name='bert_mean',
        datasource='annotation.xlsx',
        seed=0,
        max_epochs=50
    )

import torch
from huggingface_hub import hf_hub_download
from torch import nn
from treevalue import FastTreeValue

from gende import tokenize
from gende.models import load_model_from_ckpt


class ModelWithEncoding(nn.Module):
    def __init__(self, model):
        nn.Module.__init__(self)
        self.model = model

    def forward(self, x):
        encoded = self.model.encoder(x)
        predicted = self.model.ft(encoded)

        return FastTreeValue({
            'encoded': encoded,
            'predicted': predicted,
        })


_torch_stack = FastTreeValue.func(subside=True)(torch.stack)

if __name__ == '__main__':
    # use your ckpt here!
    ckpt_file = hf_hub_download(
        'HansBug/genshin-nlp-finetuning',
        'bert_mean_v1.1_seed_0/model.ckpt',
    )
    m = load_model_from_ckpt(ckpt_file)
    m = ModelWithEncoding(m)
    print(m)

    with torch.no_grad():
        # batch input
        input_ = _torch_stack([
            tokenize(
                """
        "Combat Action: When your active character is Electro Hypostasis, heal that character for 3 HP and attach the Electro Crystal Core to them.
        (You must have Electro Hypostasis in your deck to add this card to your deck.)"
                """, tokenizer_name='bert',
            ),
            tokenize(
                """
        "The character deals +1 DMG.
        When your character triggers an Elemental Reaction: Deal +1 DMG. (Twice per Round)
        (Only Catalyst Characters can equip this. A character can equip a maximum of 1 Weapon)"
                """, tokenizer_name='bert',
            ),
        ])

        print(input_)
        output = m(input_)
        print(output)