updating model loading

app.py

@@ -1,16 +1,16 @@
 import nltk
-nltk.download('wordnet')
 import spacy
-spacy.cli.download('en_core_web_sm')
+# nltk.download('wordnet')
+# spacy.cli.download('en_core_web_sm')
+
 from const import name_map
 from demo import run_gradio
-from model import EncoderDecoderVAE
+from model import get_model
 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):
@@ -24,19 +24,39 @@ def process_examples(samples, full_names):
     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'
 
-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')
+# 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)
 
-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.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)

requirements.txt

@@ -8,3 +8,4 @@ scikit-learn
 tqdm
 spacy
 sentencepiece
+lftk