adjust training procedure

first_experiment/main.py

@@ -17,23 +17,25 @@
     parser = ArgumentParser()
     parser.add_argument("--dataset", type=str, default="imdb")
     parser.add_argument("--tam_model", type=str, default="distilbert-base-uncased")
-    parser.add_argument("--embedding_model", type=str, default="distilbert-base-uncased")
+    parser.add_argument("--embedding_model", type=str, default=None)
     parser.add_argument("--init_strategy", type=str, choices=["random", "closest-to-centeroid", "furthest-to-centeroid", "expected-gradients", "certainty"], default="random")
     parser.add_argument("--stopping_criteria", type=str)
-    parser.add_argument("--dataset_size", type=int, default=[32, 64, 128, 256, 512, 1024, 2048, 4096, 8192])
+    parser.add_argument("--dataset_size", type=int, nargs="+", default=[32, 64, 128, 256, 512, 1024, 2048, 4096, 0])
     args = parser.parse_args()
 
     full_dataset = load_dataset(args.dataset)
-    #TODO stopping_critera = load_stopping_criteria(args.stopping_criteria)
     task_keys = task_to_keys[args.dataset]
 
     for dataset_size in args.dataset_size:
-        dataset = select_fewshots(
-            args,
-            full_dataset,
-            dataset_size,
-            task_keys
-        )
+        if dataset_size > 0:
+            dataset = select_fewshots(
+                args,
+                full_dataset,
+                dataset_size,
+                task_keys
+            )
+        else:
+            dataset = full_dataset
 
         train_classification(
             args,

first_experiment/selection_strategies.py

@@ -25,13 +25,17 @@ def select_fewshots(
             task_keys["label_column"]
         )
     elif args.init_strategy == "class-centeroid-closest":
+        if args.embedding_model is None:
+            raise ValueError("You need to specify an embedding model for this init strategy.")
         dataset = closest_to_centeroid_selection(
             args.embedding_model,
             full_dataset,
             dataset_size,
             task_keys
         )
     elif args.init_strategy == "class-centeroid-furthest":
+        if args.embedding_model is None:
+            raise ValueError("You need to specify an embedding model for this init strategy.")
         dataset = furthest_to_centeroid_selection(
             args.embedding_model,
             full_dataset,

first_experiment/tam_training.py

@@ -1,54 +1,47 @@
 import json
-from pathlib import Path
+from argparse import Namespace
 
 import numpy as np
 
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, DataCollatorWithPadding, TrainingArguments, Trainer
 import evaluate
+from transformers import DataCollatorWithPadding, TrainingArguments, Trainer
+
+from utils import *
 
-PATH = Path("/glusterfs/dfs-gfs-dist/goldejon/initial-starting-point-generation")
 
 def train_classification(
-    args,
-    dataset,
-    text_column,
-    label_column,
+    args: Namespace,
+    dataset: DatasetDict,
+    task_keys: dict
 ):
+    label_column = task_keys["label_column"]
     id2label = dict(enumerate(dataset["train"].features[label_column].names))
-    label2id = {v: k for k, v in id2label.items()}
 
-    model = AutoModelForSequenceClassification.from_pretrained(
-        args.tam_model,
-        num_labels=len(id2label),
-        id2label=id2label,
-        label2id=label2id
-    )
-    tokenizer = AutoTokenizer.from_pretrained(args.tam_model)
+    model, tokenizer = get_classification_model_and_tokenizer(args.tam_model, id2label=id2label)
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
-    if len(text_column) == 1:
-        def preprocess_function(examples):
-            return tokenizer(examples[text_column[0]], truncation=True)
-    else:
-        def preprocess_function(examples):
-            return tokenizer(examples[text_column[0]], examples[text_column[1]], truncation=True)
-
-    tokenized_dataset = dataset.map(preprocess_function, batched=True)
+    tokenized_dataset = dataset.map(
+        preprocess_function,
+        batched=True,
+        fn_kwargs={"tokenizer": tokenizer, "task_keys": task_keys}
+    )
 
-    if args.init_strategy in ["class-centeroid", "class-furthest"]:
-        embedding_model = args.embedding_model.split("/")[-1]
-        experiment_extension = f"{args.tam_model}_{args.dataset}_{args.dataset_size}_{args.init_strategy}_{embedding_model}"
-    else:
-        experiment_extension = f"{args.tam_model}_{args.dataset}_{args.dataset_size}_{args.init_strategy}"
+    experiment_extension = (f"{args.tam_model}"
+                            f"_{args.dataset}"
+                            f"_{args.dataset_size}"
+                            f"_{args.init_strategy}"
+                            f"_{args.embedding_model if args.embedding_model is not None else ''}")
 
     log_path = PATH / experiment_extension
 
+    batch_size = 16
+    total_steps = min(len(dataset["train"]) // batch_size * 3, 200)
     training_args = TrainingArguments(
         output_dir=str(log_path),
         learning_rate=2e-5,
         per_device_train_batch_size=16,
         per_device_eval_batch_size=16,
-        num_train_epochs=20,
+        num_train_epochs=total_steps * batch_size // len(dataset["train"]),
         warmup_ratio=0.1,
         weight_decay=0.01,
         logging_steps=5,
@@ -67,6 +60,7 @@ def compute_metrics(eval_pred):
         model=model,
         args=training_args,
         train_dataset=tokenized_dataset["train"],
+        eval_dataset=tokenized_dataset["validation"] if "validation" in tokenized_dataset else None,
         tokenizer=tokenizer,
         data_collator=data_collator,
         compute_metrics=compute_metrics,

first_experiment/utils.py

@@ -1,4 +1,4 @@
-import os
+from pathlib import Path
 from typing import Tuple
 
 import torch
@@ -7,7 +7,8 @@
 from datasets import DatasetDict
 from transformers import AutoModel, AutoModelForSequenceClassification, AutoTokenizer, PreTrainedTokenizer, PreTrainedModel
 
-CACHE_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), ".cache")
+PATH = Path("/glusterfs/dfs-gfs-dist/goldejon/initial-starting-point-generation")
+CACHE_DIR = PATH / ".cache"
 
 
 def get_embedding_model_and_tokenizer(
@@ -25,21 +26,27 @@ def get_embedding_model_and_tokenizer(
 
 def get_classification_model_and_tokenizer(
     model_name_or_path: str,
+    id2label: dict = None,
 ) -> Tuple[PreTrainedModel, PreTrainedTokenizer]:
     """
     Get classification model and tokenizer.
     """
-    model = AutoModelForSequenceClassification.from_pretrained(model_name_or_path)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        model_name_or_path,
+        num_labels=len(id2label) if id2label is not None else None,
+        id2label=id2label if id2label is not None else None,
+        label2id={label: i for i, label in id2label.items()} if id2label is not None else None,
+    )
     if torch.cuda.is_available():
         model.cuda()
     tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
     return model, tokenizer
 
+
 def preprocess_function(
     examples,
     tokenizer: PreTrainedTokenizer,
     task_keys: dict,
-    label_to_id: dict,
 ):
     sentence1_key, sentence2_key = task_keys["text_column"]
 
@@ -50,12 +57,8 @@ def preprocess_function(
     result = tokenizer(*texts, padding=True, max_length=256, truncation=True)
 
     if "label" in examples:
-        if label_to_id is not None:
-            # Map labels to IDs (not necessary for GLUE tasks)
-            result["labels"] = [label_to_id[l] for l in examples["label"]]
-        else:
-            # In all cases, rename the column to labels because the model will expect that.
-            result["labels"] = examples["label"]
+        result["labels"] = examples["label"]
+
     return result
 
 
@@ -68,20 +71,13 @@ def get_trainloader(
     """
     Get dataloader for classification dataset.
     """
-    label_column = task_keys["label_column"]
-    try:
-        label_to_id = dict(enumerate(dataset["train"].features[label_column].feature.names))
-    except (AttributeError, KeyError):
-        label_to_id = None
-
     processed_datasets = dataset.map(
         preprocess_function,
         batched=True,
         remove_columns=dataset["train"].column_names,
         fn_kwargs={
             "tokenizer": tokenizer,
             "task_keys": task_keys,
-            "label_to_id": label_to_id
         },
         desc="Running tokenizer on dataset",
     )