app.py

import nltk
nltk.download('wordnet')
import spacy
spacy.cli.download('en_core_web_sm')
from const import name_map
from demo import run_gradio
from model import EncoderDecoderVAE
from options import parse_args
import numpy as np
from transformers import T5Tokenizer
import torch
import joblib
import pandas as pd


def process_examples(samples, full_names):
    for i in range(len(samples)):
        sample = samples[i]
        input_text = tokenizer.decode(sample['sentence1_input_ids'], skip_special_tokens=True)
        ling1 = scaler.inverse_transform([sample['sentence1_ling']])[0]
        ling2 = scaler.inverse_transform([sample['sentence2_ling']])[0]
        ling = pd.DataFrame({'Index': full_names, 'Source': ling1, 'Target': ling2})
        samples[i] = [input_text, ling]
    return list(samples)

args, args_list, lng_names = parse_args(ckpt='./ckpt/model.pt')

tokenizer = T5Tokenizer.from_pretrained(args.model_name)
device = 'cuda' if torch.cuda.is_available() else 'cpu'

scaler = joblib.load('assets/scaler.bin')
full_names = [name_map[x] for x in lng_names]
samples = joblib.load('assets/samples.bin')
examples = process_examples(samples, full_names)
ling_collection = np.load('assets/ling_collection.npy')

model = EncoderDecoderVAE(args, tokenizer.pad_token_id, tokenizer.get_vocab()['</s>']).to(device)
state = torch.load(args.ckpt, map_location=torch.device('cpu'))
model.load_state_dict(state['model'], strict=False)
model.eval()

run_gradio(model, tokenizer, scaler, ling_collection, examples, full_names)