added config files and training loop code

Author: paramkpr

.gitignore

@@ -24,7 +24,8 @@ wheels/
 venv/
 env/
 ENV/
-
+venv312/
+venv311/
 # Jupyter Notebook
 .ipynb_checkpoints
 
@@ -66,4 +67,7 @@ data/clean/*
 
 # Logs
 logs/
-*.log 
+*.log 
+
+wandb/
+outputs/

.python-version

@@ -0,0 +1 @@
+3.11.2

configs/deberta_large_sst2.yaml

@@ -1,14 +1,15 @@
 model_name: microsoft/deberta-v3-large
-dataset_path: data/sst2_dd
+dataset_path: data/processed
 train_split: train
 eval_split: val
-sanity_split: sanity
+sanity_split: sent_sanity
+test_split: test
 max_len: 128
 per_device_train_batch_size: 8
 per_device_eval_batch_size: 32
 gradient_accumulation_steps: 4
 num_train_epochs: 3
-learning_rate: 2e-5
+learning_rate: 0.00005
 warmup_ratio: 0.06
 fp16: true
 logging_steps: 50

configs/deberta_large_sst2_mps.yaml

@@ -0,0 +1,20 @@
+model_name: prajjwal1/bert-small
+dataset_path: data/processed
+train_split: train
+eval_split: val
+sanity_split: sent_sanity
+test_split: test
+max_len: 128
+per_device_train_batch_size: 2
+per_device_eval_batch_size: 32
+gradient_accumulation_steps: 16
+num_train_epochs: 10
+learning_rate: 0.0005
+warmup_ratio: 0.06
+fp16: false
+logging_steps: 5
+eval_steps: 10
+save_steps: 10000
+output_dir: outputs/teacher
+report_to: wandb
+project_name: sst2_teacher_mps

requirements.txt

@@ -1,5 +1,5 @@
 torch>=1.9.0
-transformers>=4.10.0
+transformers
 datasets>=1.11.0
 numpy>=1.19.5
 scikit-learn>=0.24.2

src/train_teacher.py

@@ -29,6 +29,10 @@
 )
 logger = logging.getLogger(__name__)
 
+# Set OMP_NUM_THREADS to 1 to avoid potential CPU over-subscription
+os.environ['OMP_NUM_THREADS'] = '1'
+logger.info(f"Setting OMP_NUM_THREADS=1")
+
 
 def load_config(config_path: str) -> Dict:
     """Loads configuration from a YAML file."""
@@ -46,9 +50,11 @@ def load_datasets(dataset_path: str) -> DatasetDict:
     logger.info(f"Datasets loaded: {datasets}")
     # Ensure standard column names
     if "sentence" in datasets["train"].column_names:
-        datasets = datasets.rename_column("sentence", "text")
+        if "text" not in datasets["train"].column_names:
+            datasets = datasets.rename_column("sentence", "text")
     if "label" in datasets["train"].column_names:
-        datasets = datasets.rename_column("label", "labels")
+        if "labels" not in datasets["train"].column_names:
+            datasets = datasets.rename_column("label", "labels")
     # Make sure 'labels' column exists
     if "labels" not in datasets["train"].column_names:
         raise ValueError(
@@ -101,6 +107,7 @@ def main(config_path: str):
     # --- Setup W&B ---
     run_name = f"train_teacher_{int(time.time())}"
     if config.get("report_to") == "wandb":
+        os.environ.pop("WANDB_DISABLED", None)
         os.environ["WANDB_PROJECT"] = config["project_name"]
         logger.info(f"Logging to W&B project: {config['project_name']}")
     else:
@@ -114,7 +121,7 @@ def main(config_path: str):
     model = AutoModelForSequenceClassification.from_pretrained(
         config["model_name"], num_labels=2 # Assuming binary classification for SST-2
     )
-    tokenizer = AutoTokenizer.from_pretrained(config["model_name"], use_fast=False)
+    tokenizer = AutoTokenizer.from_pretrained(config["model_name"], use_fast=True)
 
     # --- Load and Prepare Data ---
     raw_datasets = load_datasets(config["dataset_path"])
@@ -126,8 +133,24 @@ def main(config_path: str):
     eval_dataset = tokenized_datasets[config["eval_split"]]
     sanity_dataset = tokenized_datasets[config["sanity_split"]]
 
+    ## TODO: UNCOMMENT FOR REAL RUN::
+    train_dataset = train_dataset.shuffle(seed=42).select(range(256))
+    eval_dataset = eval_dataset.shuffle(seed=42).select(range(128))
+    # sanity_dataset = sanity_dataset.shuffle(seed=42).select(range(256))
+
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
+    # --- Detect device (CUDA, MPS, or CPU) ----------------------------
+    if torch.cuda.is_available():
+        device = torch.device("cuda")
+    elif torch.backends.mps.is_available():
+        device = torch.device("mps")
+    else:
+        device = torch.device("cpu")
+
+    model.to(device)
+    logger.info(f"Using device: {device}")
+
     # --- Training Arguments ---
     logger.info("Setting up Training Arguments...")
     training_args = TrainingArguments(
@@ -144,8 +167,8 @@ def main(config_path: str):
         fp16=config["fp16"] and torch.cuda.is_available(),
         logging_dir=f"{config['output_dir']}/logs",
         logging_steps=config["logging_steps"],
-        evaluation_strategy=IntervalStrategy.STEPS,
         eval_steps=config["eval_steps"],
+        eval_strategy=IntervalStrategy.STEPS,
         save_strategy=IntervalStrategy.STEPS,
         save_steps=config["save_steps"],
         save_total_limit=2, # Saves the best and the latest checkpoints
@@ -156,6 +179,7 @@ def main(config_path: str):
         run_name=run_name,
         label_names=["labels"], # Specify label column name
         remove_unused_columns=False, # Keep all columns tokenized earlier
+        ddp_find_unused_parameters=False,
     )
     logger.info(f"FP16 enabled: {training_args.fp16}")
 
@@ -171,7 +195,7 @@ def main(config_path: str):
         model=model,
         args=training_args,
         train_dataset=train_dataset,
-        eval_dataset={"eval": eval_dataset, "sanity": sanity_dataset}, # Evaluate on both
+        eval_dataset=eval_dataset,
         tokenizer=tokenizer,
         data_collator=data_collator,
         compute_metrics=compute_metrics,
@@ -193,12 +217,7 @@ def main(config_path: str):
     trainer.log_metrics("train", metrics)
     trainer.save_metrics("train", metrics)
 
-    # Evaluate one last time on both sets with the best model
-    logger.info("Evaluating best model on eval and sanity sets...")
-    eval_metrics = trainer.evaluate(eval_dataset=eval_dataset, metric_key_prefix="final_eval")
-    trainer.log_metrics("final_eval", eval_metrics)
-    trainer.save_metrics("final_eval", eval_metrics)
-
+    # Evaluate on sanity set
     sanity_metrics = trainer.evaluate(eval_dataset=sanity_dataset, metric_key_prefix="final_sanity")
     trainer.log_metrics("final_sanity", sanity_metrics)
     trainer.save_metrics("final_sanity", sanity_metrics)

train.sh

@@ -2,7 +2,14 @@
 # Make sure to run: chmod +x train.sh
 set -e
 
-CONFIG=configs/deberta_large_sst2.yaml
+# Check if config argument was provided
+if [ $# -eq 0 ]; then
+    echo "Error: Please provide config file path as argument"
+    echo "Usage: ./train.sh <config_path>"
+    exit 1
+fi
+
+CONFIG=$1
 
 # Check if CONFIG file exists
 if [ ! -f "$CONFIG" ]; then
@@ -21,4 +28,4 @@ echo "Starting training using config: $CONFIG"
 torchrun --nnodes 1 --nproc_per_node 4 --master_port 12345 \
          src/train_teacher.py $CONFIG
 
-echo "Training script finished." 
+echo "Training script finished."