model updated

.env.example

@@ -0,0 +1 @@
+WANDB_API_KEY=<your-api-key>

README.md

@@ -1,8 +1,25 @@
 ---
 license: cc-by-nc-sa-4.0
 ---
+# CITDA:
+
+Fine-tuned `bert-base-uncased` on the `emotions` dataset
+
+Demo Notebook: https://colab.research.google.com/drive/10ZCFvlf2UV3FjU4ymf4OoipQvqHbIItG?usp=sharing
+
+## Packages
+
+- Install `torch`
+- Also, `pip install transformers datasets scikit-learn wandb seaborn python-dotenv`
+
+
+## Train
+
+1. Rename `.env.example` to `.env` and set an API key from [wandb](https://wandb.ai/authorize)
+2. You can adjust model parameters in the `explainableai.py` file.
+2. The model (`pytorch_model.bin`) is a based on the `bert-base-uncased` and already trained on the `emotions` dataset. 
+To re-produce the training run `finetune-emotions.py`. You can change the base model, or the dataset by changing that file's code.
 
-Notebook: https://colab.research.google.com/drive/10ZCFvlf2UV3FjU4ymf4OoipQvqHbIItG?usp=sharing
 
 
 ## Example

explainableai.py

@@ -3,7 +3,7 @@ from sklearn.metrics import accuracy_score, f1_score
 
 import numpy as np
 
-CITDA_EPOCHS = 10
+CITDA_EPOCHS = 6
 CITDA_WEIGHT_DECAY = 0.05 # L2 regularization
 CITDA_BATCH_SIZE = 32
 CITDA_LEARNINGRATE= 2e-5
@@ -11,7 +11,6 @@ CITDA_LEARNINGRATE= 2e-5
 class CITDA:
     def __init__(self, model, labels, base_model_name, tokenizer, encoded_data):
         self.labels = labels
-        # self.device = device
         self.tokenizer = tokenizer
         self.model = model
         self.encoded_data = encoded_data
@@ -34,25 +33,26 @@ class CITDA:
                                         weight_decay=CITDA_WEIGHT_DECAY,
                                         evaluation_strategy="epoch",
                                         save_strategy="epoch",
-                                        disable_tqdm=False)
+                                        disable_tqdm=False,
+                                        report_to="wandb")
         trainer = Trainer(model=self.model, tokenizer=self.tokenizer, args=training_args,
                     compute_metrics=compute_metrics,
                     train_dataset = self.encoded_data["train"],
-                    eval_dataset = self.encoded_data["validation"],
-                    report_to="wandb")
+                    eval_dataset = self.encoded_data["validation"])
         return trainer
 
     def train(self):
         trainer = self._get_trainer()
+        trainer.train()
         results = trainer.evaluate()
-        preds_output = trainer.predict(encoded_data["validation"])
+        preds_output = trainer.predict(self.encoded_data["validation"])
 
-        y_valid = np.array(encoded_data["validation"]["label"])
-        y_preds = np.argmax(preds_output.predictions, axis=1)
+        y_valid = np.array(self.encoded_data["validation"]["label"])
+        y_pred = np.argmax(preds_output.predictions, axis=1)
 
         #Saving the fine-tuned model
-        self.model.save_pretrained('./model') 
-        self.tokenizer.save_pretrained('./model')
+        self.model.save_pretrained('./') 
+        self.tokenizer.save_pretrained('./')
 
         return y_valid, y_pred
 

finetune-emotions.py

@@ -1,5 +1,3 @@
-# Modified https://github.com/bhadreshpsavani/ExploringSentimentalAnalysis/blob/main/SentimentalAnalysisWithDistilbert.ipynb
-
 import torch
 from sklearn.metrics import confusion_matrix
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
@@ -7,13 +5,17 @@ from datasets import load_dataset
 #import matplotlib.pyplot as plt
 import seaborn as sns
 import explainableai
+import os
+from dotenv import load_dotenv
+load_dotenv()
 
 BASE_MODEL_NAME = "bert-base-uncased"
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
 
+device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+print("Device:", device)
 
-def save_confusion_matrix(y_valid, y_preds):
-    cm = confusion_matrix(y_valid, y_preds)
+def save_confusion_matrix(y_valid, y_pred):
+    cm = confusion_matrix(y_valid, y_pred)
     f = sns.heatmap(cm, annot=True, fmt='d')
     f.figure.savefig("confusion_matrix.png")
 
@@ -24,6 +26,7 @@ def get_encoded_data(tokenizer):
     emotions_encoded = emotions.map(tokenize, batched=True, batch_size=None)
     emotions_encoded.set_format("torch", columns=["input_ids", "attention_mask", "label"])
     return emotions_encoded
+
 if __name__ == "__main__":
     labels = ['sadness', 'joy', 'love', 'anger', 'fear', 'surprise']
     model = AutoModelForSequenceClassification.from_pretrained(
@@ -38,5 +41,5 @@ if __name__ == "__main__":
     citda = explainableai.CITDA(model, labels, BASE_MODEL_NAME, tokenizer, encoded_data)
     y_valid, y_pred = citda.train()
     
-    save_confusion_matrix(y_valid, y_preds)
+    save_confusion_matrix(y_valid, y_pred)
     print("y_valid=",len(y_valid), "y_pred=", len(y_pred))

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:892b89d11f17b4c9ca7b429a72a1edab084f06a546e337090461616262c70935
-size 438018413
+oid sha256:8b0409f66d9c3fe0e3aa9976265fa1f26dbd1526cd212a696fa8d97e459b71e9
+size 438036351