From 18fe3f314873728468bb0408c00ca604bb2eaa97 Mon Sep 17 00:00:00 2001
From: Advaith Rao <advaith652@gmail.com>
Date: Mon, 4 Dec 2023 19:40:15 -0500
Subject: [PATCH] Created Base.py for base distilbert model and modified
 dp.save

---
 ethics/base.py                 | 52 ++++++++++++++++++++++++++++++++++
 ethics/differential_privacy.py | 46 +++++++++++-------------------
 2 files changed, 68 insertions(+), 30 deletions(-)
 create mode 100644 ethics/base.py

diff --git a/ethics/base.py b/ethics/base.py
new file mode 100644
index 0000000..a48536c
--- /dev/null
+++ b/ethics/base.py
@@ -0,0 +1,52 @@
+import os
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
+
+import shutil
+import pandas as pd
+import numpy as np
+from sklearn.svm import SVC
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import accuracy_score
+from sklearn.utils.class_weight import compute_class_weight
+from sklearn.pipeline import Pipeline
+import torch
+from torch import nn
+
+from transformers import DistilBertTokenizer, DistilBertForSequenceClassification, DistilBertModel
+from transformers import AdamW,get_linear_schedule_with_warmup
+
+from torch.utils.data import DataLoader, TensorDataset
+import torch.nn.functional as F
+
+import wandb
+from mlflow.sklearn import save_model
+from scipy.sparse import hstack
+
+
+class BaseDistilbertModel(nn.Module):
+    def __init__(self, num_labels, model_name='distilbert-base-uncased', device = 'cuda'):
+        super(BaseDistilbertModel, self).__init__()
+
+        # Load pre-trained RobertaModel
+        self.model = DistilBertModel.from_pretrained(model_name).to(device)
+
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        # Define classification head
+        self.classification_head = nn.Sequential(
+            nn.Linear(self.model.config.hidden_size, 128),
+            nn.ReLU(),
+            nn.Linear(128, num_labels)
+        )
+
+    def forward(self, input_ids, attention_mask, labels=None):
+        # Get model outputs
+        outputs = self.model(input_ids, attention_mask=attention_mask)
+        last_hidden_states = outputs.last_hidden_state
+
+        # Apply classification head
+        logits = self.classification_head(last_hidden_states[:, 0, :])
+
+        return logits
\ No newline at end of file
diff --git a/ethics/differential_privacy.py b/ethics/differential_privacy.py
index d21e140..0a1a6ab 100644
--- a/ethics/differential_privacy.py
+++ b/ethics/differential_privacy.py
@@ -1,6 +1,9 @@
 import os
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:512"
 
+import sys
+sys.path.append('..')
+
 import shutil
 import pandas as pd
 import numpy as np
@@ -23,39 +26,13 @@
 from mlflow.sklearn import save_model
 from scipy.sparse import hstack
 
+from base import BaseDistilbertModel
 from utils.util_modeler import Word2VecEmbedder, TPSampler
 
 from opacus import PrivacyEngine
 from opacus.utils.batch_memory_manager import BatchMemoryManager
 
 
-class BaseModel(nn.Module):
-    def __init__(self, num_labels, model_name='distilbert-base-uncased', device = 'cuda'):
-        super(BaseModel, self).__init__()
-
-        # Load pre-trained RobertaModel
-        self.model = DistilBertModel.from_pretrained(model_name).to(device)
-
-        for param in self.model.parameters():
-            param.requires_grad = False
-
-        # Define classification head
-        self.classification_head = nn.Sequential(
-            nn.Linear(self.model.config.hidden_size, 128),
-            nn.ReLU(),
-            nn.Linear(128, num_labels)
-        )
-
-    def forward(self, input_ids, attention_mask, labels=None):
-        # Get model outputs
-        outputs = self.model(input_ids, attention_mask=attention_mask)
-        last_hidden_states = outputs.last_hidden_state
-
-        # Apply classification head
-        logits = self.classification_head(last_hidden_states[:, 0, :])
-
-        return logits
-
 class DistilbertPrivacyModel:
     def __init__(
         self, 
@@ -88,7 +65,7 @@ def __init__(
         if self.path != '':
             raise NotImplementedError('Loading model from path is not implemented yet.')
         else:
-            self.model = BaseModel(num_labels=self.num_labels, model_name=self.model_name)
+            self.model = BaseDistilbertModel(num_labels=self.num_labels, model_name=self.model_name)
             self.model.to(self.device)
         
         self.privacy_engine = PrivacyEngine()
@@ -355,8 +332,17 @@ def save_model(
             os.makedirs(path, exist_ok=True)
 
         # Save the transformer model and the classification head
-        self.model.save_pretrained(path)
-        torch.save(self.classification_head.state_dict(), os.path.join(path, 'classification_head.pth'))
+        # self.model.save_pretrained(path)
+        # torch.save(self.classification_head.state_dict(), os.path.join(path, 'classification_head.pth'))
+        try:
+            torch.save(self.privacy_engine.accountant)
+        except:
+            print('Accountant not saved')
+        
+        try:
+            torch.save(self.model._module.state_dict())
+        except:
+            print('Model not saved')
     
     def accuracy(
         self,