milestone-3

.gitignore

@@ -127,4 +127,4 @@ dmypy.json
 
 # Pyre type checker
 .pyre/
-./data
+data/

milestone_3.py

@@ -1,96 +1,100 @@
-from transformers import DistilBertTokenizerFast, DistilBertModel, AdamW
+from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification, Trainer, TrainingArguments
 import torch
-from torch.utils.data import Dataset, DataLoader
+from torch.utils.data import Dataset
+# from torch.optim import AdamW
 import pandas as pd
+from sklearn.model_selection import train_test_split
 
 
 # assignment 3
 model_name = "distilbert-base-uncased"
-tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
+
+class ToxicDataset(Dataset):
+
+    def __init__(self, encodings, labels):
+        self.encodings = encodings
+        self.labels = labels
+    
+    def __getitem__(self, idx):
+        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
+        item["labels"] = torch.tensor(self.labels[idx])
+        print(item)
+        return item
+
+    def __len__(self):
+        return len(self.labels)
 
 print("Reading data...")
 data = pd.read_csv("./data/train.csv")
 toxic_data = pd.DataFrame()
 toxic_data["text"] = data["comment_text"]
 toxic_data["labels"] = data.iloc[:, 2:].values.tolist()
-print(toxic_data.head())
-
-class ToxicDataset(Dataset):
 
-    def __init__(self, dataframe, tokenizer):
-        self.tokenizer = tokenizer
-        self.data = dataframe
-        self.text = dataframe.text
-        self.labels = self.data.labels
+print("Data read. Splitting data...")
+train_texts, val_texts, train_labels, val_labels = train_test_split(toxic_data.text.to_list(), toxic_data.labels.to_list(), test_size=.2)
 
-    def __len__(self):
-        return len(self.text)
-    
-    def __getitem__(self, idx):
-        text = str(self.text[idx])
-        if len(text) > 12:
-            text = text[:12]
-
-        inputs = self.tokenizer.encode_plus(
-            text,
-            None,
-            max_length=12,
-            add_special_tokens=True,
-            pad_to_max_length=True,
-            return_token_type_ids=True
-        )
-
-        ids = inputs["input_ids"]
-        mask = inputs["attention_mask"]
-        token_type_ids = inputs["token_type_ids"]
-
-        return {
-            "ids": torch.tensor(ids, dtype=torch.long),
-            "mask": torch.tensor(mask, dtype=torch.long),
-            "token_type_ids": torch.tensor(token_type_ids, dtype=torch.long),
-            "targets": torch.tensor(self.labels[idx], dtype=torch.float)
-        }
 
+print("Data split. Tokenizing data...")
+tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
 
-print("Data read. Splitting data...")
-train_data = toxic_data.sample(frac=.8)
-test_data = toxic_data.drop(train_data.index).reset_index(drop=True)
-train_data = train_data.reset_index(drop=True)
+train_encodings = tokenizer.batch_encode_plus(train_texts, truncation=True, padding=True, return_tensors='pt')
+val_encodings = tokenizer.batch_encode_plus(val_texts, truncation=True, padding=True, return_tensors='pt')
 
-print("Data split. Tokenizing data...")
-train_set = ToxicDataset(train_data, tokenizer)
-test_set = ToxicDataset(test_data, tokenizer)
 
-train_loader = DataLoader(train_set, batch_size=8, shuffle=True, num_workers=0)
-test_loader = DataLoader(test_set, batch_size=8, shuffle=True, num_workers=0)
+train_dataset = ToxicDataset(train_encodings, train_labels)
+val_dataset = ToxicDataset(val_encodings, val_labels)
 
 print("Data tokenized. Beginning training...")
 
-device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+training_args = TrainingArguments(
+    output_dir="./results",          
+    num_train_epochs=2,             
+    per_device_train_batch_size=4, 
+    per_device_eval_batch_size=16, 
+    warmup_steps=500,               
+    weight_decay=0.01,             
+    logging_dir="./logs",            
+    logging_steps=10,
+)
+
+# device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu")
+
+model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=6)
+
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=val_dataset,
+)
+
+trainer.train()
+
+# model = DistilBertForSequenceClassification.from_pretrained(model_name, num_labels=6)
+
+# model.to(device)
+# model.train()
 
-model = DistilBertModel.from_pretrained(model_name)
-model.to(device)
-model.train()
+# train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
 
-optim = AdamW(model.parameters(), lr=5e-5)
+# optim = AdamW(model.parameters(), lr=5e-5)
 
-num_train_epochs = 2
+# num_train_epochs = 2
 
-for epoch in range(num_train_epochs):
-    for batch in train_loader:
-        optim.zero_grad()
-        input_ids = batch["ids"].to(device)
-        attention_mask = batch["mask"].to(device)
-        token_type_ids = batch["token_type_ids"].to(device, dtype = torch.long)
-        targets = batch["targets"].to(device)
+# for epoch in range(num_train_epochs):
+#     for batch in train_loader:
+#         optim.zero_grad()
+#         input_ids = batch["input_ids"].to(device)
+#         attention_mask = batch["attention_mask"].to(device)
+#         labels = batch["labels"].to(device)
 
-        outputs = model(input_ids, attention_mask, token_type_ids)
+#         outputs = model(input_ids, attention_mask=attention_mask, labels=labels)
         
-        loss = torch.nn.BCEWithLogitsLoss()(outputs, targets)
-        loss.backward()
-        optim.step()
+#         loss = outputs[0]
+#         loss.backward()
+#         optim.step()
 
-model.eval()
+# model.eval()