# -*- coding: utf-8 -*- #for some reason this import has to come first, otherwise gives segfault from transformers import BertTokenizer, BertModel import datetime import argparse import uuid import torch import random import numpy as np import pickle import os import cv2 import copy import torch.optim as optim import torch.nn as nn from tqdm import tqdm from torch.optim.lr_scheduler import StepLR from torch.utils.tensorboard import SummaryWriter from utils.util import Utils from torchtext import data from torchtext import datasets from torch.utils.data import Dataset, DataLoader from models.decoder_stgcn import Decoder from models.sent_transformers import SentiTransformerEncoder from data.dataset import SignProdDataset import spacy import nltk from torch import nn import torch.nn.functional as fnn from torch.autograd import Variable from torch.optim import SGD from torchvision.datasets import LSUN from torchvision import transforms from torch.utils.data import Dataset from torchvision.utils import make_grid PAD_IDX = 0 tokenizer = BertTokenizer.from_pretrained("bert-base-uncased") max_input_length = tokenizer.max_model_input_sizes['bert-base-uncased'] nlp = spacy.load("en_core_web_sm") def tokenize_and_cut(sentence): tokens = tokenizer.tokenize(sentence) tokens = tokens[:max_input_length-2] return tokens def get_pos(text): doc = nlp(text) pos_list = list() for token in doc: pos_list.append(token.pos_) return pos_list def get_dataset(data_root, batch_size, device, train = True): init_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.cls_token) eos_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.sep_token) pad_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.pad_token) unk_token_idx = tokenizer.convert_tokens_to_ids(tokenizer.unk_token) TEXT = data.Field( batch_first = True, use_vocab = False, tokenize = tokenize_and_cut, preprocessing = tokenizer.convert_tokens_to_ids, init_token = init_token_idx, eos_token = eos_token_idx, pad_token = pad_token_idx, unk_token = unk_token_idx ) RAW = data.RawField() dataset = SignProdDataset( data_root, train, fields = [('src', TEXT), ('kps', RAW), ("s_feature", RAW), ("pos", RAW), ("aus", RAW), ("z", RAW), ("idx", RAW), ("file_name", RAW)] ) iterator = data.BucketIterator( dataset = dataset, batch_size=batch_size, device=device, sort_key=lambda x: len(x.src), repeat=False, sort=False, shuffle=True if train else False, sort_within_batch=True ) return iterator def save_Zs(output_root, Z): with open(output_root, "wb") as handler: pickle.dump(Z, handler) def read_pickle(file_fp): with open(file_fp, "rb") as handler: return pickle.load(handler) def read_sentences(dataset_root, files): sentences = list() for idx, file_ in enumerate(files[:50]): file_fp = os.path.join(dataset_root, file_) instance = read_pickle(file_fp) sentences.append((idx, instance["text"])) return sentences def get_k_similar_sentences(pos, sentences, k): sentences_dist = dict() pos_rev_dict = build_pos_rev_dict() for idx, sentence in sentences: pos_s = get_pos(sentence) pos_s_mapped = [pos_rev_dict[p] for p in pos_s] pos_mapped = [pos_rev_dict[p] for p in pos] pos_s_mapped = "".join(pos_s_mapped) pos_mapped = "".join(pos_mapped) distance = nltk.edit_distance(pos_s_mapped, pos_mapped) if distance not in sentences_dist: sentences_dist[distance] = [idx] else: sentences_dist[distance].append(idx) sorted_keys = list(sorted(sentences_dist.keys())) selected_sentences = list() for key in sorted_keys: sentences = sentences_dist[key] for s in sentences: if len(selected_sentences) == k: return selected_sentences selected_sentences.append(s) def build_pos_rev_dict(): return dict( ADJ = "a", ADP = "b", ADV = "c", AUX = "d", CONJ = "e", CCONJ = "f", DET = "g", INTJ = "h", NOUN = "i", NUM = "j", PART = "k", PRON = "l", PROPN = "m", PUNCT = "n", SCONJ = "o", SYM = "p", VERB = "q", X = "r", SPACE = "s" ) def project_l2_ball(z): """ project the vectors in z onto the l2 unit norm ball""" return z / np.maximum(np.sqrt(np.sum(z**2, axis=1))[:, np.newaxis], 1) def main(args): SEED = 1234 random.seed(SEED) np.random.seed(SEED) torch.manual_seed(SEED) torch.backends.cudnn.deterministic = True device = torch.device('cuda:{}'.format(args.device)) if args.device != -1 else torch.device('cuda') N_EPOCHS = args.epochs dataset_root = args.dataset_root #encoder = SentiTransformerEncoder().to(device) decoder = Decoder(device).to(device) decoder.load_state_dict(torch.load(args.decoder_ckpt)) decoder.eval() print("Reading Zs") Z = read_pickle(args.zs_ckpt) print("Reading Zsent") Zs = read_pickle(args.zsent_ckpt) """ test_iterator = get_dataset( args.dataset_root, args.batch_size, device, train = False ) """ files = sorted(os.listdir(dataset_root)) print("Reading sentences") sentences = read_sentences(dataset_root, files) for file_ in files: file_fp = os.path.join(dataset_root, file_) instance = read_pickle(file_fp) kps = instance["kps"] text = "It cannot be turned around this is the easy way out" pos = get_pos(text) selected_sentences = get_k_similar_sentences(pos, sentences, 2) z = 0 zs = 0 for idx_s in selected_sentences: z += Z[idx_s] zs += Zs[idx_s] z = np.expand_dims(z, axis = 0) zs = np.expand_dims(zs, axis = 0) z = project_l2_ball(z) zs = project_l2_ball(zs) with torch.no_grad(): zs = torch.Tensor(zs).to(device) z = torch.Tensor(z).to(device) input_ = z + zs faces_fake = decoder(None, None, input_).permute(0, 2, 3, 1).cpu().numpy() Utils.visualize_data_single(faces_fake, "test_test3.mp4") import pdb pdb.set_trace() if __name__ == "__main__": parser = argparse.ArgumentParser(description='Process some integers.') parser.add_argument('--device', '-dev', type = int, default = 0) parser.add_argument('--batch_size', '-bs', type = int, default = 42) parser.add_argument('--epochs', '-e', type = int, default = 200000) parser.add_argument('--dataset_root', '-d', type = str, default = "/srv/storage/datasets/rafaelvieira/slp_dataset_f68_full_ps_new_arch_sample_wz_final_grammar/train") parser.add_argument('--outputs_root', '-ot', type = str, default = "/srv/storage/datasets/rafaelvieira/text2expression/results/") parser.add_argument('--decoder_ckpt', '-dckpt', type = str, default = "/srv/storage/datasets/rafaelvieira/text2expression/step_162500/decoder.pth") parser.add_argument('--zs_ckpt', '-zckpt', type = str, default = "/srv/storage/datasets/rafaelvieira/text2expression/step_162500/Zs.pkl") parser.add_argument('--zsent_ckpt', '-zsentckpt', type = str, default = "/srv/storage/datasets/rafaelvieira/text2expression/step_162500/Zsent.pkl") args = parser.parse_args() main(args)