import torch
from transformers import AutoModel, AutoTokenizer, FlaxAutoModel
from datasets import load_dataset
from wechsel import WECHSEL, load_embeddings

source_tokenizer = AutoTokenizer.from_pretrained("roberta-large")
model = AutoModel.from_pretrained("roberta-large")

target_tokenizer = AutoTokenizer.from_pretrained("./")

wechsel = WECHSEL(
    load_embeddings("en"),
    load_embeddings("fi"),
    bilingual_dictionary="finnish"
)

target_embeddings, info = wechsel.apply(
    source_tokenizer,
    target_tokenizer,
    model.get_input_embeddings().weight.detach().numpy(),
)

model.get_input_embeddings().weight.data = torch.from_numpy(target_embeddings)

model.save_pretrained("./")

flax_model = FlaxAutoModel.from_pretrained("./", from_pt=True)
flax_model.save_pretrained("./")