def test_data(self):
path_dir_name = os.path.dirname(os.path.realpath(__file__))
data_path = os.path.join(path_dir_name, "sample.txt")
with tempfile.TemporaryDirectory() as tmpdirname:
processor = TextProcessor()
processor.train_tokenizer([data_path],
vocab_size=1000,
to_save_dir=tmpdirname,
languages={
"<mzn>": 0,
"<glk": 1
})
create_batches.write(text_processor=processor,
cache_dir=tmpdirname,
seq_len=512,
txt_file=data_path,
sen_block_size=10)
dataset = TextDataset(save_cache_dir=tmpdirname, max_cache_size=3)
assert dataset.line_num == 70
dataset.__getitem__(3)
assert len(dataset.current_cache) == 3
dataset.__getitem__(9)
assert len(dataset.current_cache) == 3
dataset.__getitem__(69)
assert len(dataset.current_cache) == 2
def get_tokenizer(tokenizer_path: Optional[str] = None,
train_path: Optional[str] = None,
model_path: Optional[str] = None,
vocab_size: Optional[int] = None) -> TextProcessor:
if tokenizer_path is None:
if not os.path.exists(model_path):
os.makedirs(model_path)
print("Training Tokenizer...")
text_processor = TextProcessor()
print("Writing raw text...")
languages = set()
with open(train_path + ".tmp", "w") as wf:
with open(train_path, "r") as rf:
for i, line in enumerate(rf):
spl = [
sen.strip() for sen in line.split("</s>")
if len(sen.strip()) > 0
]
if len(spl) == 0: continue
if spl[0].startswith("<"):
sen_split = spl[0].strip().split(" ")
spl[0] = " ".join(sen_split[1:])
languages.add(sen_split[0])
wf.write("\n".join(spl))
wf.write("\n")
if ((i + 1) % 1000 == 0):
print(i + 1, "\r", end="")
print("Writing raw text done!")
text_processor.train_tokenizer(
paths=[train_path + ".tmp"],
vocab_size=vocab_size,
to_save_dir=model_path,
languages={l: i
for i, l in enumerate(sorted(languages))})
print("Removing temporary file!")
os.system("rm " + train_path + ".tmp &")
print("done!")
else:
text_processor = TextProcessor(tokenizer_path)
return text_processor
def load(out_dir: str):
text_processor = TextProcessor(tok_model_path=out_dir)
with open(os.path.join(out_dir, "config"), "rb") as fp:
config = pickle.load(fp)
if isinstance(config, dict):
# For older configs
config = BertConfig(**config)
lm = LM(text_processor=text_processor, config=config)
lm.load_state_dict(
torch.load(os.path.join(out_dir, "model.state_dict")))
return lm
def test_train_tokenizer(self):
path_dir_name = os.path.dirname(os.path.realpath(__file__))
data_path = os.path.join(path_dir_name, "sample.txt")
with tempfile.TemporaryDirectory() as tmpdirname:
processor = TextProcessor()
processor.train_tokenizer([data_path],
vocab_size=1000,
to_save_dir=tmpdirname,
languages={"<en>": 0})
assert processor.tokenizer.get_vocab_size() == 1000
sen1 = "Obama signed many landmark bills into law during his first two years in office."
assert processor._tokenize(sen1) is not None
many_sens = "\n".join([sen1] * 10)
assert len(processor.tokenize(many_sens)) == 10
new_prcoessor = TextProcessor(tok_model_path=tmpdirname)
assert new_prcoessor.tokenizer.get_vocab_size() == 1000
sen2 = "Obama signed many landmark bills into law during his first two years in office."
assert processor._tokenize(sen2) is not None
def train(options):
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
num_processors = max(torch.cuda.device_count(), 1)
mt_model = SenSim(text_processor=text_processor, enc_layer=options.encoder_layer, embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
if options.pretrained_path is not None:
pret = Seq2Seq.load(Seq2Seq, options.pretrained_path, tok_dir=options.tokenizer_path)
mt_model.init_from_lm(pret)
print("Model initialization done!")
optimizer = build_optimizer(mt_model, options.learning_rate, warump_steps=options.warmup)
trainer = SenSimTrainer(model=mt_model, mask_prob=options.mask_prob, optimizer=optimizer, clip=options.clip,
fp16=options.fp16)
pin_memory = torch.cuda.is_available()
mt_train_loader = SenSimTrainer.get_mt_train_data(mt_model, num_processors, options, pin_memory)
src_neg_data = dataset.MassDataset(batch_pickle_dir=options.src_neg,
max_batch_capacity=num_processors * options.total_capacity * 5,
max_batch=num_processors * options.batch * 5,
pad_idx=mt_model.text_processor.pad_token_id(), keep_pad_idx=False,
max_seq_len=options.max_seq_len, keep_examples=False)
dst_neg_data = dataset.MassDataset(batch_pickle_dir=options.dst_neg,
max_batch_capacity=num_processors * options.total_capacity * 5,
max_batch=num_processors * options.batch * 5,
pad_idx=mt_model.text_processor.pad_token_id(), keep_pad_idx=False,
max_seq_len=options.max_seq_len, keep_examples=False)
src_neg_loader = data_utils.DataLoader(src_neg_data, batch_size=1, shuffle=True, pin_memory=pin_memory)
dst_neg_loader = data_utils.DataLoader(dst_neg_data, batch_size=1, shuffle=True, pin_memory=pin_memory)
mt_dev_loader = None
if options.mt_dev_path is not None:
mt_dev_loader = SenSimTrainer.get_mt_dev_data(mt_model, options, pin_memory, text_processor, trainer, )
step, train_epoch = 0, 1
trainer.best_loss = 1000000
while options.step > 0 and step < options.step:
print("train epoch", train_epoch)
step = trainer.train_epoch(mt_train_iter=mt_train_loader, max_step=options.step, mt_dev_iter=mt_dev_loader,
saving_path=options.model_path, step=step, src_neg_iter=src_neg_loader,
dst_neg_iter=dst_neg_loader)
train_epoch += 1
def load(out_dir: str, tok_dir: str):
text_processor = TextProcessor(tok_model_path=tok_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
with open(os.path.join(out_dir, "mt_config"), "rb") as fp:
enc_layer, embed_dim, intermediate_dim = pickle.load(fp)
model = SenSim(text_processor=text_processor,
enc_layer=enc_layer,
embed_dim=embed_dim,
intermediate_dim=intermediate_dim)
model.load_state_dict(torch.load(os.path.join(
out_dir, "mt_model.state_dict"),
map_location=device),
strict=False)
return model, text_processor
def train(options):
lex_dict = None
if options.dict_path is not None:
lex_dict = get_lex_dict(options.dict_path)
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
image_captioner = Seq2Seq.load(ImageCaptioning, options.pretrained_path, tok_dir=options.tokenizer_path)
txt2ImageModel = Caption2Image(text_processor=text_processor, enc_layer=options.encoder_layer,
embed_dim=options.embed_dim, intermediate_dim=options.intermediate_layer_dim)
print("Model initialization done!")
# We assume that the collator function returns a list with the size of number of gpus (in case of cpus,
collator = dataset.ImageTextCollator()
num_batches = max(1, torch.cuda.device_count())
optimizer = build_optimizer(txt2ImageModel, options.learning_rate, warump_steps=options.warmup)
trainer = Caption2ImageTrainer(model=txt2ImageModel, caption_model=image_captioner, mask_prob=options.mask_prob,
optimizer=optimizer,
clip=options.clip,
beam_width=options.beam_width, max_len_a=options.max_len_a,
max_len_b=options.max_len_b, len_penalty_ratio=options.len_penalty_ratio,
fp16=options.fp16, mm_mode=options.mm_mode)
pin_memory = torch.cuda.is_available()
img_train_loader = ImageMTTrainer.get_img_loader(collator, dataset.ImageCaptionDataset, options.train_path,
txt2ImageModel, num_batches, options, pin_memory,
lex_dict=lex_dict)
img_dev_loader = ImageMTTrainer.get_img_loader(collator, dataset.ImageCaptionDataset, options.dev_path,
txt2ImageModel, num_batches, options, pin_memory,
lex_dict=lex_dict,
shuffle=False, denom=2)
step, train_epoch = 0, 1
while options.step > 0 and step < options.step:
print("train epoch", train_epoch)
step = trainer.train_epoch(img_data_iter=img_train_loader, img_dev_data_iter=img_dev_loader,
max_step=options.step, lex_dict=lex_dict,
saving_path=options.model_path, step=step)
train_epoch += 1
def test_albert_init(self):
path_dir_name = os.path.dirname(os.path.realpath(__file__))
data_path = os.path.join(path_dir_name, "sample.txt")
with tempfile.TemporaryDirectory() as tmpdirname:
processor = TextProcessor()
processor.train_tokenizer([data_path],
vocab_size=1000,
to_save_dir=tmpdirname,
languages={"<en>": 0})
lm = LM(text_processor=processor)
assert lm.encoder.base_model.embeddings.word_embeddings.num_embeddings == 1000
lm.save(tmpdirname)
new_lm = LM.load(tmpdirname)
assert new_lm.config == lm.config
def test_albert_seq2seq_init(self):
path_dir_name = os.path.dirname(os.path.realpath(__file__))
data_path = os.path.join(path_dir_name, "sample.txt")
with tempfile.TemporaryDirectory() as tmpdirname:
processor = TextProcessor()
processor.train_tokenizer([data_path],
vocab_size=1000,
to_save_dir=tmpdirname,
languages={
"<en>": 0,
"<fa>": 1
})
seq2seq = Seq2Seq(text_processor=processor)
src_inputs = torch.tensor([[
1, 2, 3, 4, 5,
processor.pad_token_id(),
processor.pad_token_id()
], [1, 2, 3, 4, 5, 6, processor.pad_token_id()]])
tgt_inputs = torch.tensor(
[[6, 8, 7,
processor.pad_token_id(),
processor.pad_token_id()],
[6, 8, 7, 8, processor.pad_token_id()]])
src_mask = (src_inputs != processor.pad_token_id())
tgt_mask = (tgt_inputs != processor.pad_token_id())
src_langs = torch.tensor([[0], [0]]).squeeze()
tgt_langs = torch.tensor([[1], [1]]).squeeze()
seq_output = seq2seq(src_inputs,
tgt_inputs,
src_mask,
tgt_mask,
src_langs,
tgt_langs,
log_softmax=True)
assert list(seq_output.size()) == [5, processor.vocab_size()]
seq_output = seq2seq(src_inputs, tgt_inputs, src_mask, tgt_mask,
src_langs, tgt_langs)
assert list(seq_output.size()) == [5, processor.vocab_size()]
def load(cls, out_dir: str, tok_dir: str):
text_processor = TextProcessor(tok_model_path=tok_dir)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
with open(os.path.join(out_dir, "mt_config"), "rb") as fp:
lang_dec, use_proposals, enc_layer, dec_layer, embed_dim, intermediate_dim, tie_embed, resnet_depth, freeze_image = pickle.load(
fp)
mt_model = cls(text_processor=text_processor,
lang_dec=lang_dec,
use_proposals=use_proposals,
tie_embed=tie_embed,
enc_layer=enc_layer,
dec_layer=dec_layer,
embed_dim=embed_dim,
intermediate_dim=intermediate_dim,
freeze_image=freeze_image,
resnet_depth=resnet_depth)
mt_model.load_state_dict(torch.load(os.path.join(
out_dir, "mt_model.state_dict"),
map_location=device),
strict=False)
return mt_model
parser.add_option(
"--all",
action="store_true",
dest="use_all",
help="Choose all sentences instead of only subset of captions",
default=False)
parser.add_option("--only",
action="store_true",
dest="only",
help="Choose only captions",
default=False)
(options, args) = parser.parse_args()
return options
if __name__ == "__main__":
options = get_options()
tokenizer = TextProcessor(options.tokenizer_path)
print("Writing batches")
write(text_processor=tokenizer,
output_file=options.output_file,
input_file=options.file,
root_img_dir=options.image_dir,
skip_check=options.skip_check,
max_len=options.max_len,
ref_file=options.ref,
only_captions=options.only,
choose_relevant=not options.use_all)
print("Finished")
def save_data(data, already_queried, session_token, checkpoint_number, data_type, save_path, args):
all_data = {}
if checkpoint_number == 8:
already_queried = [] # Because we start with a different dataset at checkpoint 9
metadata = {}
metadata['already_queried'] = already_queried
metadata['session_token'] = session_token
metadata['checkpoint_number'] = checkpoint_number
all_data["metadata"] = metadata
tok_path = os.path.dirname(os.path.realpath(__file__)) + "/tok"
if data_type == "train":
np.save(os.path.join(save_path, "metadata.npy"), all_data)
eng = []
ar = []
for element in data:
if element["english"] and element["arabic"]:
cleaned_eng = element["english"].replace("\r", "").replace("\n"," ")
cleaned_ar = element["arabic"].replace("\r", "").replace("\n"," ")
if len(cleaned_ar) > 0 and len(cleaned_eng) > 0:
eng.append(cleaned_eng)
ar.append(cleaned_ar)
if checkpoint_number != 1:
ar_prev, eng_prev = load_training_data(save_path)
ar = ar_prev + ar
eng = eng_prev + eng
save_to_file(eng, save_path, "english.train")
save_to_file(ar, save_path, "arabic.train")
tokenizer = TextProcessor(tok_path)
create_mt_batches.write(text_processor=tokenizer, output_file=os.path.join(save_path, "train.batch"),
src_txt_file=os.path.join(save_path, "arabic.train"),
dst_txt_file=os.path.join(save_path, "english.train"))
train_options = TrainOptions()
train_options.mt_train_path = os.path.join(save_path, "train.batch")
num_iters = max(10, (len(eng) / (train_options.batch / 100)) * 200)
train_options.step = int(min(args["iter"], num_iters))
print("Training for", train_options.step, "iterations!")
train_options.model_path = os.path.join(save_path, "train.model")
train_options.tokenizer_path = tok_path
train_options.pretrained_path = os.path.join(save_path, "pret.model")
train_options.encoder_layer = args["enc"]
train_options.decoder_layer = args["dec"]
train_options.embed_dim = args["embed"]
train_options.beam_width = args["beam"]
train_image_mt.ImageMTTrainer.train(train_options)
print("Training Done!")
elif data_type == "test":
ar = []
ID = []
for key in data.keys():
ID.append(key)
ar.append(data[key])
save_to_file(ar, save_path, "arabic.test")
save_to_file(ID, save_path, "ID.test")
print("Translating ...")
translate_options = TranslateOptions()
translate_options.mt_train_path = os.path.join(save_path, "train.batch")
translate_options.model_path = os.path.join(save_path, "train.model")
translate_options.tokenizer_path = tok_path
translate_options.input_path = os.path.join(save_path, "arabic.test")
translate_options.output_path = os.path.join(save_path, "english.test.output")
translate_options.beam_width = args["beam"]
translate.translate(translate_options)
print("Translating done!")
def train(options):
lex_dict = None
if options.dict_path is not None:
lex_dict = get_lex_dict(options.dict_path)
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
num_processors = max(torch.cuda.device_count(), 1)
if options.pretrained_path is not None:
print("Loading pretrained path", options.pretrained_path)
mt_model = Seq2Seq.load(ImageMassSeq2Seq,
options.pretrained_path,
tok_dir=options.tokenizer_path)
else:
mt_model = ImageMassSeq2Seq(
use_proposals=lex_dict is not None,
tie_embed=options.tie_embed,
text_processor=text_processor,
resnet_depth=options.resnet_depth,
lang_dec=options.lang_decoder,
enc_layer=options.encoder_layer,
dec_layer=options.decoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
if options.lm_path is not None:
lm = LM(text_processor=text_processor,
enc_layer=options.encoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
mt_model.init_from_lm(lm)
print("Model initialization done!")
# We assume that the collator function returns a list with the size of number of gpus (in case of cpus,
collator = dataset.ImageTextCollator()
num_batches = max(1, torch.cuda.device_count())
if options.continue_train:
with open(os.path.join(options.pretrained_path, "optim"),
"rb") as fp:
optimizer = pickle.load(fp)
else:
optimizer = build_optimizer(mt_model,
options.learning_rate,
warump_steps=options.warmup)
trainer = ImageMTTrainer(model=mt_model,
mask_prob=options.mask_prob,
optimizer=optimizer,
clip=options.clip,
beam_width=options.beam_width,
max_len_a=options.max_len_a,
max_len_b=options.max_len_b,
len_penalty_ratio=options.len_penalty_ratio,
fp16=options.fp16,
mm_mode=options.mm_mode)
pin_memory = torch.cuda.is_available()
mt_train_loader = None
if options.mt_train_path is not None:
mt_train_loader = ImageMTTrainer.get_mt_train_data(
mt_model,
num_processors,
options,
pin_memory,
lex_dict=lex_dict)
mt_dev_loader = None
if options.mt_dev_path is not None:
mt_dev_loader = ImageMTTrainer.get_mt_dev_data(mt_model,
options,
pin_memory,
text_processor,
trainer,
lex_dict=lex_dict)
step, train_epoch = 0, 1
while options.step > 0 and step < options.step and train_epoch <= 10:
print("train epoch", train_epoch, "step:", step)
step = trainer.train_epoch(mt_train_iter=mt_train_loader,
max_step=options.step,
lex_dict=lex_dict,
mt_dev_iter=mt_dev_loader,
saving_path=options.model_path,
step=step,
save_opt=False)
train_epoch += 1
def train(options):
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
lm_class = ReformerLM if options.reformer else LM
if options.pretrained_path is None:
lm = lm_class(text_processor=text_processor,
size=options.model_size)
else:
lm = lm_class.load(options.pretrained_path)
if options.reformer:
lm.config.hidden_dropout_prob = options.dropout
lm.config.local_attention_probs_dropout_prob = options.dropout
lm.config.lsh_attention_probs_dropout_prob = options.dropout
else:
LMTrainer.config_dropout(lm, options.dropout)
train_data = dataset.TextDataset(save_cache_dir=options.train_path,
max_cache_size=options.cache_size)
dev_data = dataset.TextDataset(save_cache_dir=options.dev_path,
max_cache_size=options.cache_size,
load_all=True)
if options.continue_train:
with open(os.path.join(options.pretrained_path, "optim"),
"rb") as fp:
optimizer = pickle.load(fp)
else:
optimizer = build_optimizer(lm, options.learning_rate,
options.warmup)
trainer = LMTrainer(model=lm,
mask_prob=options.mask_prob,
optimizer=optimizer,
clip=options.clip)
collator = dataset.TextCollator(pad_idx=text_processor.pad_token_id())
train_sampler, dev_sampler = None, None
pin_memory = torch.cuda.is_available()
loader = data_utils.DataLoader(train_data,
batch_size=options.batch,
shuffle=False,
pin_memory=pin_memory,
collate_fn=collator,
sampler=train_sampler)
dev_loader = data_utils.DataLoader(dev_data,
batch_size=options.batch,
shuffle=False,
pin_memory=pin_memory,
collate_fn=collator,
sampler=dev_sampler)
step, train_epoch = 0, 1
while step <= options.step:
print("train epoch", train_epoch)
step = trainer.train_epoch(data_iter=loader,
dev_data_iter=dev_loader,
saving_path=options.model_path,
step=step)
def train(options):
lex_dict = None
if options.dict_path is not None:
lex_dict = get_lex_dict(options.dict_path)
if not os.path.exists(options.model_path):
os.makedirs(options.model_path)
text_processor = TextProcessor(options.tokenizer_path)
assert text_processor.pad_token_id() == 0
if options.pretrained_path is not None:
mt_model = Seq2Seq.load(ImageCaptioning,
options.pretrained_path,
tok_dir=options.tokenizer_path)
else:
mt_model = ImageCaptioning(
use_proposals=lex_dict is not None,
tie_embed=options.tie_embed,
text_processor=text_processor,
resnet_depth=options.resnet_depth,
lang_dec=options.lang_decoder,
enc_layer=options.encoder_layer,
dec_layer=options.decoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
if options.lm_path is not None:
lm = LM(text_processor=text_processor,
enc_layer=options.encoder_layer,
embed_dim=options.embed_dim,
intermediate_dim=options.intermediate_layer_dim)
mt_model.init_from_lm(lm)
print("Model initialization done!")
# We assume that the collator function returns a list with the size of number of gpus (in case of cpus,
collator = dataset.ImageTextCollator()
num_batches = max(1, torch.cuda.device_count())
if options.continue_train:
with open(os.path.join(options.pretrained_path, "optim"),
"rb") as fp:
optimizer = pickle.load(fp)
else:
optimizer = build_optimizer(mt_model,
options.learning_rate,
warump_steps=options.warmup)
trainer = ImageCaptionTrainer(
model=mt_model,
mask_prob=options.mask_prob,
optimizer=optimizer,
clip=options.clip,
beam_width=options.beam_width,
max_len_a=options.max_len_a,
max_len_b=options.max_len_b,
len_penalty_ratio=options.len_penalty_ratio,
fp16=options.fp16,
mm_mode=options.mm_mode)
pin_memory = torch.cuda.is_available()
img_train_loader = ImageMTTrainer.get_img_loader(
collator,
dataset.ImageCaptionDataset,
options.train_path,
mt_model,
num_batches,
options,
pin_memory,
lex_dict=lex_dict)
img_dev_loader = ImageMTTrainer.get_img_loader(
collator,
dataset.ImageCaptionDataset,
options.dev_path,
mt_model,
num_batches,
options,
pin_memory,
lex_dict=lex_dict,
shuffle=False,
denom=2)
trainer.reference = None
if img_dev_loader is not None:
trainer.reference = []
generator = (trainer.generator.module if hasattr(
trainer.generator, "module") else trainer.generator)
for data in img_dev_loader:
for batch in data:
captions = [b["captions"] for b in batch]
for caption in captions:
refs = get_outputs_until_eos(
text_processor.sep_token_id(),
caption,
remove_first_token=True)
ref = [
generator.seq2seq_model.text_processor.tokenizer.
decode(ref.numpy()) for ref in refs
]
trainer.reference += ref
step, train_epoch = 0, 1
while options.step > 0 and step < options.step:
print("train epoch", train_epoch)
step = trainer.train_epoch(img_data_iter=img_train_loader,
img_dev_data_iter=img_dev_loader,
max_step=options.step,
lex_dict=lex_dict,
saving_path=options.model_path,
step=step)
train_epoch += 1