Update DistilBERT.py

DistilBERT.py

@@ -11,6 +11,7 @@ import pandas as pd
 import numpy as np
 
 # Điều chỉnh các tham số
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 MAX_LEN = 100
 TRAIN_BATCH_SIZE = 4
 VALID_BATCH_SIZE = 4
@@ -19,7 +20,11 @@ LEARNING_RATE = 1e-05
 tokenizer_DB = DistilBertTokenizer.from_pretrained('distilbert-base-uncased', truncation=True, do_lower_case=True)
 
 # Tạo dataframe
-
+train_df_DB = pd.read_csv('./data/train.csv')
+train_df_DB['label'] = train_df_DB.iloc[:, 1:].values.tolist()
+test_df_DB = pd.read_csv('./data/test.csv')
+test_df_DB = test_df_DB[['text', 'preprocess_sentence', 'label']]
+test_df_DB['label'] = test_df_DB.iloc[:, 2:].values.tolist()
 
 # Tạo class
 class BinaryLabel(Dataset):
@@ -72,4 +77,66 @@ training_set = BinaryLabel(train_df_DB, tokenizer, MAX_LEN)
 testing_set = BinaryLabel(test_df_DB, tokenizer, MAX_LEN)
 
 training_loader = DataLoader(training_set, **train_params)
-testing_loader = DataLoader(testing_set, **test_params)
+testing_loader = DataLoader(testing_set, **test_params)
+
+# Create model
+class DistilBERTClass(torch.nn.Module):
+    def __init__(self):
+        super(DistilBERTClass, self).__init__()
+        self.l1 = DistilBertModel.from_pretrained("distilbert-base-uncased")
+        self.pre_classifier = torch.nn.Linear(768, 768)
+        self.dropout = torch.nn.Dropout(0.1)
+        self.classifier = torch.nn.Linear(768, 1)
+
+    def forward(self, input_ids, attention_mask, token_type_ids):
+        output_1 = self.l1(input_ids=input_ids, attention_mask=attention_mask)
+        hidden_state = output_1[0]
+        pooler = hidden_state[:, 0]
+        pooler = self.pre_classifier(pooler)
+        pooler = torch.nn.ReLU()(pooler)
+        pooler = self.dropout(pooler)
+        output = self.classifier(pooler)
+        return output
+
+model_DB = DistilBERTClass()
+model_DB.to(device)
+
+# Validation function
+def validation(testing_loader):
+    model_DB.eval()
+    fin_targets=[]
+    fin_outputs=[]
+    with torch.no_grad():
+        for _, data in tqdm(enumerate(testing_loader, 0)):
+            ids = data['ids'].to(device, dtype = torch.long)
+            mask = data['mask'].to(device, dtype = torch.long)
+            token_type_ids = data['token_type_ids'].to(device, dtype = torch.long)
+            targets = data['targets'].to(device, dtype = torch.float)
+            outputs = model_DB(ids, mask, token_type_ids)
+            fin_targets.extend(targets.cpu().detach().numpy().tolist())
+            fin_outputs.extend(torch.sigmoid(outputs).cpu().detach().numpy().tolist())
+    return fin_outputs, fin_targets
+
+# Train function
+def train(epoch):
+    model.train()
+    for _,data in tqdm(enumerate(training_loader, 0)):
+        ids = data['ids'].to(device, dtype = torch.long)
+        mask = data['mask'].to(device, dtype = torch.long)
+        token_type_ids = data['token_type_ids'].to(device, dtype = torch.long)
+        targets = data['targets'].to(device, dtype = torch.float)
+
+        outputs = model(ids, mask, token_type_ids)
+
+        optimizer.zero_grad()
+        loss = loss_fn(outputs, targets)
+        if _%50==0:
+            print(f'Epoch: {epoch}, Loss:  {loss.item()}')
+        if loss.item() < 0.07:
+            print(f'Breaking the loop as loss is below 0.07: {loss.item()}')
+            break
+        loss.backward()
+        optimizer.step()
+
+for epoch in range(3):
+    train(epoch)