Delete train_script.py

train_script.py

@@ -1,164 +0,0 @@
-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 vietnamese