app.py

import nltk
import spacy
# nltk.download('wordnet')
# spacy.cli.download('en_core_web_sm')

from const import name_map
from demo import run_gradio
from model import get_model
from options import parse_args
import numpy as np
from transformers import T5Tokenizer
import torch
import joblib


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')
print(args)
exit()

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

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')

scaler = joblib.load('assets/scaler.bin')
model, ling_disc, sem_emb = get_model(args, tokenizer, device)

state = torch.load(args.ckpt, map_location=torch.device('cpu'))
model.load_state_dict(state['model'], strict=True)
model.eval()
print(model is not None, ling_disc is not None, sem_emb is not None)
exit()

if args.disc_type == 't5':
    state = torch.load(args.disc_ckpt)
    if 'model' in state:
        ling_disc.load_state_dict(state['model'], strict=False)
    else:
        ling_disc.load_state_dict(state, strict=False)
ling_disc.eval()

state = torch.load(args.sem_ckpt)
if 'model' in state:
    sem_emb.load_state_dict(state['model'], strict=False)
else:
    sem_emb.load_state_dict(state, strict=False)
sem_emb.eval()

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