DistilBert_RR.py

# -*- coding: utf-8 -*-
"""distilBERT_for_legal_dataset_finall.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1DWq6VUHp68DcJ3rYqZBGt8_YHqp8-G0u
"""

import numpy as np
import torch
import torch.nn as nn
import torchmetrics

from transformers import DistilBertTokenizer, DistilBertModel
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
dbert_model = DistilBertModel.from_pretrained("distilbert-base-uncased")

X_train=np.load('texts_train.npy')
X_test=np.load('texts_dev.npy')
y_train=np.load('labels_train.npy')
y_test=np.load('labels_dev.npy')

X_train.shape,X_test.shape,y_train.shape,y_test.shape

tokenized_text_train=[]
tokenized_text_test=[]
for text in X_train:
  encoded_input = tokenizer.encode_plus(text,max_length=32,padding='max_length', truncation=True,return_tensors="pt")
  tokenized_text_train.append(encoded_input)
for text in X_test:
  encoded_input = tokenizer.encode_plus(text,max_length=32,padding='max_length',truncation=True, return_tensors="pt")
  tokenized_text_test.append(encoded_input)

num_classes=13
classifier=nn.Sequential(
    nn.Linear(768,512),
    nn.ReLU(),
    nn.Dropout(0.4),
    nn.Linear(512,128),
    nn.ReLU(),
    nn.Dropout(0.4),
    nn.Linear(128,num_classes),
    nn.Softmax()
)

import torch
class custom(nn.Module):
  def __init__(self):
    super().__init__()
    self.m1=dbert_model
    self.m2=classifier
  def forward(self,inp,mask):
    x=self.m1(input_ids=inp,attention_mask=mask).last_hidden_state
    x=torch.squeeze(x, 0)
    x= torch.mean(x,1)
    x=self.m2(x)
    return x

db_model = custom()

tokenized_text_train[0]

inp_train=[]
masks_train=[]
tti_train=[]
for data in tokenized_text_train:
  inp_train.append(data['input_ids'])
  masks_train.append(data['attention_mask'])


inp_test=[]
masks_test=[]
tti_test=[]
for data in tokenized_text_test:
  inp_test.append(data['input_ids'])
  masks_test.append(data['attention_mask'])


inp_train=torch.stack(inp_train)
masks_train=torch.stack(masks_train)


inp_test=torch.stack(inp_test)
masks_test=torch.stack(masks_test)

label_mapping = {
    'ANALYSIS': 0,
    'ARG_PETITIONER': 1,
    'ARG_RESPONDENT': 2,
    'FAC': 3,
    'ISSUE': 4,
    'NONE': 5,
    'PREAMBLE': 6,
    'PRE_NOT_RELIED': 7,
    'PRE_RELIED': 8,
    'RATIO': 9,
    'RLC': 10,
    'RPC': 11,
    'STA': 12
}

y_train=torch.LongTensor(y_train)
y_test=torch.LongTensor(y_test)

from torch.utils.data import DataLoader,TensorDataset

train_dataset=TensorDataset(inp_train,masks_train,y_train)
val_dataset=TensorDataset(inp_test,masks_test,y_test)

train = DataLoader(train_dataset, batch_size=64, shuffle=True)
val=DataLoader(val_dataset, batch_size=64, shuffle=True)

epochs=5

criterion=nn.CrossEntropyLoss()
opt=torch.optim.Adam(db_model.parameters(), lr=0.1)
# accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes).to('cuda')
train_accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
val_accuracy = torchmetrics.Accuracy(task="multiclass", num_classes=num_classes)
for epoch in range(epochs):
  avg_train_acc=0
  avg_val_acc=0
  count=0
  avg_train_loss=0
  avg_val_loss=0

  for batch in train:
    count=count+1
    print(f'Epoch {epoch} Batch no.: {count}')
    X_batch_in,X_batch_mask,label_batch = batch

    X_batch_in=torch.squeeze(X_batch_in)
    X_batch_mask=torch.squeeze(X_batch_mask)

    preds=db_model(X_batch_in,X_batch_mask)

    loss=criterion(preds,label_batch)
    acc=train_accuracy.update(preds,label_batch)
    # avg_train_acc=avg_train_acc+acc
    avg_train_loss=avg_train_loss+loss

    opt.zero_grad()
    loss.backward()
    opt.step()

  with torch.no_grad():
    for batch in val:

      X_val_batch_in,X_val_batch_mask,label_val_batch = batch
      X_val_batch_in=torch.squeeze(X_val_batch_in)

      X_val_batch_mask=torch.squeeze(X_val_batch_mask)


      val_preds=db_model(X_val_batch_in,X_val_batch_mask)

      val_loss=criterion(val_preds,label_val_batch)
      val_acc=val_accuracy.update(val_preds,label_val_batch)
      # avg_val_acc=avg_val_acc+val_acc
      avg_val_loss=avg_val_loss+val_loss
  print(f"| Epoch={epoch} | Training Accuracy={train_accuracy.compute()} | Validation Accuracy={val_accuracy.compute()} | Training Loss={avg_train_loss/len(train)} | Validation_Loss={avg_val_loss/len(val)} |")
  print('-------------------------------------------------------------------------------------------------------------------------------------------------------------------------')

db_model=db_model.to('cpu')
torch.save(db_model,'test_db_model.pt')