refactoring notebook code

src/model_training.py

@@ -0,0 +1,120 @@
+from datasets import DatasetDict, Dataset
+import numpy as np
+import pandas as pd
+import torch
+from sklearn.linear_model import LogisticRegression
+from sklearn.metrics import accuracy_score, f1_score
+from transformers import AutoModel, AutoTokenizer
+
+from .utils import serialize_data, load_data
+
+
+class PreProcessor:
+    def __init__(self, model_name, train_path:str, test_path:str, output_path:str):
+        self.device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+        self.model = AutoModel.from_pretrained(model_name).to(self.device)
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.df_train = pd.read_csv(train_path, sep="\t")
+        self.df_test = pd.read_csv(test_path, sep="\t")
+        self.output_path = output_path
+
+    def _get_datasetdict_object(self):
+        mapper = {"#2_tweet": "tweet", "#3_country_label": "label"}
+        columns_to_keep = ["tweet", "label"]
+
+        df_train = self.df_train.rename(columns=mapper)[columns_to_keep]
+        df_test = self.df_test.rename(columns=mapper)[columns_to_keep]
+
+        train_dataset = Dataset.from_pandas(df_train)
+        test_dataset = Dataset.from_pandas(df_test)
+        data = DatasetDict({'train': train_dataset, 'test': test_dataset})
+
+        return data
+    
+    def _tokenize(self, batch):
+        return self.tokenizer(batch["tweet"], padding=True)
+
+    def _encode_data(self, data):
+        data_encoded = data.map(self._tokenize, batched=True, batch_size=None)
+        return data_encoded
+
+    def _extract_hidden_states(self, batch):
+        inputs = {k:v.to(self.device) for k,v in batch.items()
+                    if k in self.tokenizer.model_input_names}
+        with torch.no_grad():
+            last_hidden_state = self.model(**inputs).last_hidden_state
+
+        return {"hidden_state": last_hidden_state[:,0].cpu().numpy()}
+    
+    def _get_features(self, data_encoded):
+        data_encoded.set_format("torch", columns=["input_ids", "attention_mask", "label"])
+        data_hidden = data_encoded.map(self._extract_hidden_states, batched=True, batch_size=50)
+        return data_hidden  
+        
+    def preprocess_data(self):
+        data = self._get_datasetdict_object()
+        data_encoded = self._encode_data(data)
+        data_hidden = self._get_features(data_encoded)
+        serialize_data(data_hidden, output_path=self.output_path)
+
+
+class Model():
+    def __init__(self, data_input_path:str, model_name:str):
+        self.model_name = model_name
+        self.model = None
+        self.data = load_data(input_path=data_input_path)
+        self.X_train = np.array(self.data["train"]["hidden_state"])
+        self.X_test = np.array(self.data["test"]["hidden_state"])
+        self.y_train = np.array(self.data["train"]["label"])
+        self.y_test = np.array(self.data["test"]["label"])
+
+    def _train_logistic_regression(X_train, y_train):
+        lr_model = LogisticRegression(multi_class='multinomial', 
+                                    class_weight="balanced", 
+                                    max_iter=1000, 
+                                    random_state=2024)
+        lr_model.fit(X_train, y_train)
+        return lr_model
+
+    def train_model(self, output_path):
+        if self.model_name != "lr":
+            raise ValueError(f"Model name {self.model_name} does not exist. Please try 'lr'!")
+
+        lr_model = self._train_logistic_regression(self.X_train, self.y_train)
+        self.model = lr_model
+        serialize_data(lr_model, output_path)
+    
+    def _get_metrics(self, y_true, y_preds):
+        accuracy = accuracy_score(y_true, y_preds)
+        f1_macro = f1_score(y_true, y_preds, average="macro")
+        f1_weighted = f1_score(y_true, y_preds, average="weighted")
+        print(f"Accuracy: {accuracy}")
+        print(f"F1 macro average: {f1_macro}")
+        print(f"F1 weighted average: {f1_weighted}")
+
+    def evaluate_predictions(self):
+        train_preds = self.model.predict(self.X_train)
+        test_preds = self.model.predict(self.X_test)
+
+        print(self.model_name)
+        print("\nTrain set:")
+        self._get_metrics(self.y_train, train_preds)
+        print("-"*50)
+        print("Test set:")
+        self._get_metrics(self.y_test, test_preds)
+
+
+def main():
+    file_path = "../data/data_hidden.pkl"
+    preprocessor = PreProcessor(model_name="moussaKam/AraBART",
+                                train_path="../data/DA_train_labeled.tsv",
+                                test_path="../data/DA_dev_labeled.tsv",
+                                output_path=file_path)
+    preprocessor.preprocess_data()
+    model = Model(data_input_path=file_path, model_name="lr")
+    model.train_model("../models/logistic_regression.pkl")
+    model.evaluate_predictions()
+
+if __name__ == "__main__":
+    main()
+