add trainer

source/services/predicting_effective_arguments/train/model.py

@@ -5,6 +5,7 @@ import torch
 import numpy as np
 import torch.nn.functional as F
 import matplotlib.pyplot as plt
+from typing import List
 from sklearn.metrics import ConfusionMatrixDisplay, confusion_matrix
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -22,35 +23,37 @@ class TransformersSequenceClassifier:
         self.tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)
         self.model = AutoModelForSequenceClassification.from_pretrained(model_checkpoint, num_labels=num_labels).to(device)
 
-    def tokenizer_func(self, batch):
+    def tokenizer_batch(self, batch):
         return self.tokenizer(batch["inputs"], truncation=True) #, max_len=386
 
-    def train(self, train_dataset, eval_dataset, epochs=2, batch_size=64):
-        
-        train_tok_dataset = train_dataset.map(self.tokenizer_func, batched=True, remove_columns=('inputs', '__index_level_0__'))
-        val_tok_dataset = eval_dataset.map(self.tokenizer_func, batched=True, remove_columns=('inputs', '__index_level_0__'))
+    def tokenize_dataset(self, dataset):
+        return dataset.map(self.tokenizer_batch, batched=True, remove_columns=('inputs', '__index_level_0__'))
+    def train(self, train_dataset, eval_dataset, batch_size, epochs):
         data_collator = DataCollatorWithPadding(tokenizer=self.tokenizer, padding='longest')
-        training_args = TrainingArguments(output_dir=f"{self.model_output_dir}/predicting_effective_arguments_distilbert",
+        training_args = TrainingArguments(output_dir=self.model_output_dir,
                                           num_train_epochs=epochs,
                                           learning_rate=2e-5,
                                           per_device_train_batch_size=batch_size,
                                           per_device_eval_batch_size=batch_size,
                                           weight_decay=0.01,
                                           evaluation_strategy="epoch",
+                                          save_strategy='epoch',
                                           disable_tqdm=False,
                                           logging_steps=len(train_dataset)// batch_size,
                                           push_to_hub=True,
+                                          load_best_model_at_end=True,
                                           log_level="error")
         self.trainer = Trainer(
             model=self.model,
             args=training_args,
             compute_metrics=self._compute_metrics,
-            train_dataset=train_tok_dataset,
-            eval_dataset=val_tok_dataset,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
             tokenizer=self.tokenizer,
             data_collator=data_collator
         )
         self.trainer.train()
+        self.trainer.push_to_hub(commit_message="Training completed!")
 
     @staticmethod
     def _compute_metrics(pred):
@@ -96,9 +99,9 @@ class TransformersSequenceClassifier:
         return y_preds
     
     @staticmethod
-    def predict_test_data(model_checkpoint, test_data):
+    def predict_test_data(model_checkpoint, test_list: List[str]) -> List:
         pipe_classifier = pipeline("text-classification", model=model_checkpoint)
-        preds = pipe_classifier(test_data, return_all_scores=True)
+        preds = pipe_classifier(test_list, return_all_scores=True)
         return preds
 
 

source/services/predicting_effective_arguments/train/seq_classification.py

@@ -7,15 +7,11 @@ from datasets import Dataset, load_metric
 from sklearn.model_selection import train_test_split
 from source.services.predicting_effective_arguments.train.model import TransformersSequenceClassifier
 
-TARGET = 'discourse_effectiveness'
-TEXT = "discourse_text"
-MODEL_CHECKPOINT = "distilbert-base-uncased"
-MODEL_OUTPUT_DIR ='source/services/predicting_effective_arguments/model/hf_textclassification'
 class CFG:
     TARGET = 'discourse_effectiveness'
     TEXT = "discourse_text"
     MODEL_CHECKPOINT = "distilbert-base-uncased"
-    MODEL_OUTPUT_DIR ='source/services/predicting_effective_arguments/model/hf_textclassification'
+    MODEL_OUTPUT_DIR ='source/services/predicting_effective_arguments/model/hf_textclassification/predicting_effective_arguments_distilbert'
     model_name="debertav3base"
     learning_rate=1.5e-5
     weight_decay=0.02
@@ -28,7 +24,6 @@ class CFG:
     save_steps=100
     max_length=512
 
-tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)
 
 def seed_everything(seed: int):
     import random, os
@@ -52,17 +47,37 @@ def prepare_input_text(df, sep_token):
 if __name__ == '__main__':
 
     config = CFG()
+    tokenizer = AutoTokenizer.from_pretrained(config.MODEL_CHECKPOINT)
     seqClassifer = TransformersSequenceClassifier(model_output_dir=config.MODEL_OUTPUT_DIR, tokenizer=tokenizer, model_checkpoint="distilbert-base-uncased", num_labels=3) #distilbert-base-uncased
     data = pd.read_csv("data/raw_data/train.csv")[:100]
     test_df = pd.read_csv("data/raw_data/test.csv")
-    train_df, valid_df = train_test_split(data, test_size=0.30, random_state=42)
+    train_size = 0.7
+    valid_size = 0.2
+    test_size = 0.1
+
+    # First split: Separate out the training set
+    train_df, temp_df = train_test_split(data, test_size=1 - train_size)
+
+    # Second split: Separate out the validation and test sets
+    valid_df, test_df = train_test_split(temp_df, test_size=test_size / (test_size + valid_size))
+
 
     train_df = prepare_input_text(train_df, sep_token=tokenizer.sep_token)
     valid_df = prepare_input_text(valid_df, sep_token=tokenizer.sep_token)
+    test_df = prepare_input_text(test_df, sep_token=tokenizer.sep_token)
     
-    train_dataset = Dataset.from_pandas(train_df[['inputs', TARGET]]).rename_column(TARGET, 'label').class_encode_column("label")
-    val_dataset = Dataset.from_pandas(valid_df[['inputs', TARGET]]).rename_column(TARGET, 'label').class_encode_column("label")
-    seqClassifer.train(train_dataset=train_dataset, eval_dataset=val_dataset, epochs=2, batch_size=64)
+    train_dataset = Dataset.from_pandas(train_df[['inputs', config.TARGET]]).rename_column(config.TARGET, 'label').class_encode_column("label")
+    val_dataset = Dataset.from_pandas(valid_df[['inputs', config.TARGET]]).rename_column(config.TARGET, 'label').class_encode_column("label")
+    test_dataset = Dataset.from_pandas(test_df[['inputs', config.TARGET]]).rename_column(config.TARGET, 'label').class_encode_column("label")
+
+    train_tok_dataset = seqClassifer.tokenize_dataset(dataset=train_dataset)
+    val_tok_dataset = seqClassifer.tokenize_dataset(dataset=val_dataset)
+    test_tok_dataset = seqClassifer.tokenize_dataset(dataset=test_dataset)
+
+    seqClassifer.train(train_dataset=train_tok_dataset, eval_dataset=val_tok_dataset, epochs=1, batch_size=16)
+    y_pred = seqClassifer.predict_valid_data(val_tok_dataset)
+    seqClassifer.predict_test_data(model_checkpoint=config.MODEL_OUTPUT_DIR, test_data=test_df['inputs'].tolist())
+    pass
 
     """
     train_df[TARGET].value_counts(ascending=True).plot.barh()