upload

aec7e95 over 4 years ago

4.92 kB

import argparse
import logging
from torch.utils.data import Dataset, IterableDataset
import gzip
import json
from transformers import Seq2SeqTrainer, AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainingArguments
import sys
from datetime import datetime
import torch
import random
from shutil import copyfile
import os
import wandb
import random
import re
from datasets import load_dataset
import tqdm


logging.basicConfig(
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S",
    handlers=[logging.StreamHandler(sys.stdout)],
)

parser = argparse.ArgumentParser()
parser.add_argument("--lang", required=True)
parser.add_argument("--model_name", default="google/mt5-base")
parser.add_argument("--epochs", default=4, type=int)
parser.add_argument("--batch_size", default=32, type=int)
parser.add_argument("--max_source_length", default=320, type=int)
parser.add_argument("--max_target_length", default=64, type=int)
parser.add_argument("--eval_size", default=1000, type=int)
#parser.add_argument("--fp16", default=False, action='store_true')
args = parser.parse_args()

wandb.init(project="doc2query", name=f"{args.lang}-{args.model_name}")





def main():
    ############ Load dataset
    queries = {}
    for row in tqdm.tqdm(load_dataset('unicamp-dl/mmarco', f'queries-{args.lang}')['train']):
        queries[row['id']] = row['text']

    """
    collection = {}
    for row in tqdm.tqdm(load_dataset('unicamp-dl/mmarco', f'collection-{args.lang}')['collection']):
        collection[row['id']] = row['text']
    """
    collection = load_dataset('unicamp-dl/mmarco', f'collection-{args.lang}')['collection']

    train_pairs = []
    eval_pairs = []


    with open('qrels.train.tsv') as fIn:
        for line in fIn:
            qid, _, did, _ = line.strip().split("\t")

            qid = int(qid)
            did = int(did)

            assert did == collection[did]['id']
            text = collection[did]['text']

            pair = (queries[qid], text)
            if len(eval_pairs) < args.eval_size:
                eval_pairs.append(pair)
            else:
                train_pairs.append(pair)

  
    print(f"Train pairs: {len(train_pairs)}")


    ############ Model
    model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name)
    tokenizer = AutoTokenizer.from_pretrained(args.model_name)

    save_steps = 1000

    output_dir = 'output/'+args.lang+'-'+args.model_name.replace("/", "-")+'-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    print("Output dir:", output_dir)

    # Write self to path
    os.makedirs(output_dir, exist_ok=True)

    train_script_path = os.path.join(output_dir, 'train_script.py')
    copyfile(__file__, train_script_path)
    with open(train_script_path, 'a') as fOut:
        fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))

    ####

    training_args = Seq2SeqTrainingArguments(
        output_dir=output_dir,
        bf16=True,
        per_device_train_batch_size=args.batch_size,
        evaluation_strategy="steps",
        save_steps=save_steps,
        logging_steps=100,
        eval_steps=save_steps, #logging_steps,
        warmup_steps=1000,
        save_total_limit=1,
        num_train_epochs=args.epochs,
        report_to="wandb",
    )

    ############ Arguments

    ############ Load datasets


    print("Input:", train_pairs[0][1])
    print("Target:", train_pairs[0][0])

    print("Input:", eval_pairs[0][1])
    print("Target:", eval_pairs[0][0])


    def data_collator(examples):
        targets = [row[0] for row in examples]
        inputs = [row[1] for row in examples]
        label_pad_token_id = -100

        model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8 if training_args.fp16 else None)

        # Setup the tokenizer for targets
        with tokenizer.as_target_tokenizer():
            labels = tokenizer(targets, max_length=args.max_target_length, padding=True, truncation=True, pad_to_multiple_of=8 if training_args.fp16 else None)

        # replace all tokenizer.pad_token_id in the labels by -100 to ignore padding in the loss.
        labels["input_ids"] = [
            [(l if l != tokenizer.pad_token_id else label_pad_token_id) for l in label] for label in labels["input_ids"]
        ]


        model_inputs["labels"] = torch.tensor(labels["input_ids"])
        return model_inputs

    ## Define the trainer
    trainer = Seq2SeqTrainer(
        model=model,
        args=training_args,
        train_dataset=train_pairs,
        eval_dataset=eval_pairs,
        tokenizer=tokenizer,
        data_collator=data_collator
    )

    ### Save the model
    train_result = trainer.train()
    trainer.save_model()
    
    
if __name__ == "__main__":
    main()

# Script was called via:
#python train_hf_trainer_multilingual.py --lang arabic

	import argparse
	import logging
	from torch.utils.data import Dataset, IterableDataset
	import gzip
	import json
	from transformers import Seq2SeqTrainer, AutoModelForSeq2SeqLM, AutoTokenizer, Seq2SeqTrainingArguments
	import sys
	from datetime import datetime
	import torch
	import random
	from shutil import copyfile
	import os
	import wandb
	import random
	import re
	from datasets import load_dataset
	import tqdm


	logging.basicConfig(
	format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
	datefmt="%Y-%m-%d %H:%M:%S",
	handlers=[logging.StreamHandler(sys.stdout)],
	)

	parser = argparse.ArgumentParser()
	parser.add_argument("--lang", required=True)
	parser.add_argument("--model_name", default="google/mt5-base")
	parser.add_argument("--epochs", default=4, type=int)
	parser.add_argument("--batch_size", default=32, type=int)
	parser.add_argument("--max_source_length", default=320, type=int)
	parser.add_argument("--max_target_length", default=64, type=int)
	parser.add_argument("--eval_size", default=1000, type=int)
	#parser.add_argument("--fp16", default=False, action='store_true')
	args = parser.parse_args()

	wandb.init(project="doc2query", name=f"{args.lang}-{args.model_name}")





	def main():
	############ Load dataset
	queries = {}
	for row in tqdm.tqdm(load_dataset('unicamp-dl/mmarco', f'queries-{args.lang}')['train']):
	queries[row['id']] = row['text']

	"""
	collection = {}
	for row in tqdm.tqdm(load_dataset('unicamp-dl/mmarco', f'collection-{args.lang}')['collection']):
	collection[row['id']] = row['text']
	"""
	collection = load_dataset('unicamp-dl/mmarco', f'collection-{args.lang}')['collection']

	train_pairs = []
	eval_pairs = []


	with open('qrels.train.tsv') as fIn:
	for line in fIn:
	qid, _, did, _ = line.strip().split("\t")

	qid = int(qid)
	did = int(did)

	assert did == collection[did]['id']
	text = collection[did]['text']

	pair = (queries[qid], text)
	if len(eval_pairs) < args.eval_size:
	eval_pairs.append(pair)
	else:
	train_pairs.append(pair)


	print(f"Train pairs: {len(train_pairs)}")


	############ Model
	model = AutoModelForSeq2SeqLM.from_pretrained(args.model_name)
	tokenizer = AutoTokenizer.from_pretrained(args.model_name)

	save_steps = 1000

	output_dir = 'output/'+args.lang+'-'+args.model_name.replace("/", "-")+'-'+datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
	print("Output dir:", output_dir)

	# Write self to path
	os.makedirs(output_dir, exist_ok=True)

	train_script_path = os.path.join(output_dir, 'train_script.py')
	copyfile(__file__, train_script_path)
	with open(train_script_path, 'a') as fOut:
	fOut.write("\n\n# Script was called via:\n#python " + " ".join(sys.argv))

	####

	training_args = Seq2SeqTrainingArguments(
	output_dir=output_dir,
	bf16=True,
	per_device_train_batch_size=args.batch_size,
	evaluation_strategy="steps",
	save_steps=save_steps,
	logging_steps=100,
	eval_steps=save_steps, #logging_steps,
	warmup_steps=1000,
	save_total_limit=1,
	num_train_epochs=args.epochs,
	report_to="wandb",
	)

	############ Arguments

	############ Load datasets


	print("Input:", train_pairs[0][1])
	print("Target:", train_pairs[0][0])

	print("Input:", eval_pairs[0][1])
	print("Target:", eval_pairs[0][0])


	def data_collator(examples):
	targets = [row[0] for row in examples]
	inputs = [row[1] for row in examples]
	label_pad_token_id = -100

	model_inputs = tokenizer(inputs, max_length=args.max_source_length, padding=True, truncation=True, return_tensors='pt', pad_to_multiple_of=8 if training_args.fp16 else None)

	# Setup the tokenizer for targets
	with tokenizer.as_target_tokenizer():
	labels = tokenizer(targets, max_length=args.max_target_length, padding=True, truncation=True, pad_to_multiple_of=8 if training_args.fp16 else None)

	# replace all tokenizer.pad_token_id in the labels by -100 to ignore padding in the loss.
	labels["input_ids"] = [
	[(l if l != tokenizer.pad_token_id else label_pad_token_id) for l in label] for label in labels["input_ids"]
	]


	model_inputs["labels"] = torch.tensor(labels["input_ids"])
	return model_inputs

	## Define the trainer
	trainer = Seq2SeqTrainer(
	model=model,
	args=training_args,
	train_dataset=train_pairs,
	eval_dataset=eval_pairs,
	tokenizer=tokenizer,
	data_collator=data_collator
	)

	### Save the model
	train_result = trainer.train()
	trainer.save_model()


	if __name__ == "__main__":
	main()

	# Script was called via:
	#python train_hf_trainer_multilingual.py --lang arabic