diff --git a/generate.py b/generate.py
index 405912d..ef617ed 100644
--- a/generate.py
+++ b/generate.py
@@ -9,48 +9,50 @@ import numpy as np
 from model import BERT_Arch
 from sklearn.metrics import classification_report
 from sklearn.metrics import accuracy_score
+from transformers import BertTokenizerFast, BertForSequenceClassification
+from transformers import Trainer, TrainingArguments
 
-path = 'saved_weights.pt'
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-device = torch.device("cuda")
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self, encodings, labels):
+        self.encodings = encodings
+        self.labels = labels
 
-bert = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')
+    def __getitem__(self, idx):
+        item = {k: torch.tensor(v[idx]) for k, v in self.encodings.items()}
+        item["labels"] = torch.tensor([self.labels[idx]])
+        return item
 
-model = BERT_Arch(bert)
-model.load_state_dict(torch.load(path))
-model.to(device)
+    def __len__(self):
+        return len(self.labels)
 
-test_data = pd.read_csv("dev-0/in.tsv", sep="\t")
-test_data.columns = ["text", "d"]
+device = torch.device('cuda')
 
-test_target = pd.read_csv("dev-0/expected.tsv", sep="\t")
+train_texts = \
+    pd.read_csv('train/in.tsv.xz', compression='xz', sep='\t', header=None, error_bad_lines=False, quoting=3)[0].tolist()[:1000]
+train_labels = pd.read_csv('train/expected.tsv', sep='\t', header=None, quoting=3)[0].tolist()[:1000]
+dev_texts = pd.read_csv('dev-0/in.tsv.xz', compression='xz', sep='\t', header=None, quoting=3)[0].tolist()[:1000]
+dev_labels = pd.read_csv('dev-0/expected.tsv', sep='\t', header=None, quoting=3)[0].tolist()[:1000]
+model_name = "bert-base-uncased"
 
-tokens_train = tokenizer.batch_encode_plus(
-    test_data["text"],
-    max_length = 25,
-    padding='max_length',
-    truncation=True
-)
+model = BertForSequenceClassification.from_pretrained(
+    model_name, num_labels=len(pd.unique(train_labels))).to(device)
+max_length = 512
+tokenizer = BertTokenizerFast.from_pretrained(model_name, do_lower_case=True)
 
-test_seq = torch.tensor(tokens_train['input_ids'])
-test_mask = torch.tensor(tokens_train['attention_mask'])
+# model.load_pretrained(model_path)
+# tokenizer.load_pretrainded(model_path)
 
-#define a batch size
-batch_size = 32
+train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=max_length)
+valid_encodings = tokenizer(dev_texts, truncation=True, padding=True, max_length=max_length)
 
-# wrap tensors
-test_data = TensorDataset(test_seq, test_mask)
-
-# sampler for sampling the data during training
-test_sampler = RandomSampler(test_data)
-
-# dataLoader for train set
-test_dataloader = DataLoader(test_data, sampler=test_sampler, batch_size=batch_size)
+input_ids = torch.tensor(valid_encodings.data['input_ids'])[:100]
+attention_mask = torch.tensor(valid_encodings.data['attention_mask'])[:100]
 
 with torch.no_grad():
-  preds = model(test_seq.to(device), test_mask.to(device))
-  preds = preds.detach().cpu().numpy()
+  preds = model(input_ids.to(device), attention_mask.to(device))
+  preds = preds.logits.detach().cpu().numpy()
   preds = np.argmax(preds, axis = 1)
     
-print(classification_report(test_target['0'], preds))
-print(accuracy_score(test_target['0'], preds))
\ No newline at end of file
+print(preds)
+print(classification_report(dev_labels, preds))
+print(accuracy_score(dev_labels, preds))
\ No newline at end of file
diff --git a/main.py b/main.py
index 655375b..76754b0 100644
--- a/main.py
+++ b/main.py
@@ -1,215 +1,90 @@
-import pandas as pd
-from transformers import BertTokenizer, AdamW, AutoModelForSequenceClassification
 import torch
-import matplotlib.pyplot as plt
-from torch.utils.data import TensorDataset, DataLoader, RandomSampler
-import torch.nn as nn
-from sklearn.utils.class_weight import compute_class_weight
+from transformers.file_utils import is_torch_available
+from transformers import BertTokenizerFast, BertForSequenceClassification
+from transformers import Trainer, TrainingArguments
 import numpy as np
-from model import BERT_Arch
-
-train_input_path = "train/in.tsv"
-train_target_path = "train/expected.tsv"
-
-train_input = pd.read_csv(train_input_path, sep="\t")
-train_input.columns=["text", "d"]
-train_target = pd.read_csv(train_target_path, sep="\t")
-
-tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-device = torch.device("cuda")
+import random
+from sklearn.metrics import accuracy_score
+import pandas as pd
 
 
-# seq_len = [len(i.split()) for i in train_input["text"]]
+class Dataset(torch.utils.data.Dataset):
+    def __init__(self, encodings, labels):
+        self.encodings = encodings
+        self.labels = labels
 
-# pd.Series(seq_len).hist(bins = 30)
-# plt.show()
+    def __getitem__(self, idx):
+        item = {k: torch.tensor(v[idx]) for k, v in self.encodings.items()}
+        item["labels"] = torch.tensor([self.labels[idx]])
+        return item
 
-bert = AutoModelForSequenceClassification.from_pretrained('bert-base-uncased')
+    def __len__(self):
+        return len(self.labels)
 
 
-tokens_train = tokenizer.batch_encode_plus(
-    train_input["text"],
-    max_length = 25,
-    padding='max_length',
-    truncation=True
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    if is_torch_available():
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+def compute_metrics(pred):
+    labels = pred.label_ids
+    preds = pred.predictions.argmax(-1)
+    acc = accuracy_score(labels, preds)
+    return {
+        'accuracy': acc,
+    }
+
+set_seed(1)
+
+train_texts = \
+    pd.read_csv('train/in.tsv.xz', compression='xz', sep='\t', header=None, error_bad_lines=False, quoting=3)[0].tolist()
+train_labels = pd.read_csv('train/expected.tsv', sep='\t', header=None, quoting=3)[0].tolist()
+dev_texts = pd.read_csv('dev-0/in.tsv.xz', compression='xz', sep='\t', header=None, quoting=3)[0].tolist()
+dev_labels = pd.read_csv('dev-0/expected.tsv', sep='\t', header=None, quoting=3)[0].tolist()
+# test_texts = pd.read_table('test-A/in.tsv.xz', compression='xz', sep='\t', header=None, quoting=3)
+
+model_name = "bert-base-uncased"
+max_length = 25
+tokenizer = BertTokenizerFast.from_pretrained(model_name, do_lower_case=True)
+
+train_encodings = tokenizer(train_texts, truncation=True, padding=True, max_length=max_length)
+valid_encodings = tokenizer(dev_texts, truncation=True, padding=True, max_length=max_length)
+
+train_dataset = Dataset(train_encodings, train_labels)
+valid_dataset = Dataset(valid_encodings, dev_labels)
+
+model = BertForSequenceClassification.from_pretrained(
+    model_name, num_labels=len(pd.unique(train_labels))).to("cuda")
+
+training_args = TrainingArguments(
+    output_dir='./results',          # output directory
+    num_train_epochs=3,              # total number of training epochs
+    per_device_train_batch_size=1,  # batch size per device during training
+    per_device_eval_batch_size=1,   # batch size for evaluation
+    warmup_steps=500,                # number of warmup steps for learning rate scheduler
+    weight_decay=0.01,               # strength of weight decay
+    logging_dir='./logs',            # directory for storing logs
+    load_best_model_at_end=True,     # load the best model when finished training (default metric is loss)
+    # but you can specify `metric_for_best_model` argument to change to accuracy or other metric
+    logging_steps=200,               # log & save weights each logging_steps
+    evaluation_strategy="steps",     # evaluate each `logging_steps`
 )
 
-train_seq = torch.tensor(tokens_train['input_ids'])
-train_mask = torch.tensor(tokens_train['attention_mask'])
-train_y = torch.tensor(train_target.to_numpy())
+trainer = Trainer(
+    model=model,                         # the instantiated Transformers model to be trained
+    args=training_args,                  # training arguments, defined above
+    train_dataset=train_dataset,         # training dataset
+    eval_dataset=valid_dataset,          # evaluation dataset
+    compute_metrics=compute_metrics,     # the callback that computes metrics of interest
+)
 
-#define a batch size
-batch_size = 32
+trainer.train()
 
-# wrap tensors
-train_data = TensorDataset(train_seq, train_mask, train_y)
+trainer.evaluate()
 
-# sampler for sampling the data during training
-train_sampler = RandomSampler(train_data)
-
-# dataLoader for train set
-train_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=batch_size)
-
-for param in bert.parameters():
-    param.requires_grad = False
-
-model = BERT_Arch(bert)
-model = model.to(device)
-# model.cuda(0)
-
-optimizer = AdamW(model.parameters(), lr = 1e-5)
-
-class_weights = compute_class_weight('balanced', np.unique(train_target.to_numpy()), train_target['1'])
-weights= torch.tensor(class_weights,dtype=torch.float)
-weights = weights.to(device)
-
-# define the loss function
-cross_entropy  = nn.NLLLoss(weight=weights) 
-
-# number of training epochs
-epochs = 10
-
-def train():
-  
-  model.train()
-
-  total_loss, total_accuracy = 0, 0
-  
-  # empty list to save model predictions
-  total_preds=[]
-  
-  # iterate over batches
-  for step,batch in enumerate(train_dataloader):
-    
-    # progress update after every 50 batches.
-    if step % 50 == 0 and not step == 0:
-      print('  Batch {:>5,}  of  {:>5,}.'.format(step, len(train_dataloader)))
-
-    # push the batch to gpu
-    batch = [r.to(device) for r in batch]
- 
-    sent_id, mask, labels = batch
-
-    # clear previously calculated gradients 
-    model.zero_grad()        
-
-    # get model predictions for the current batch
-    preds = model(sent_id, mask)
-
-    # compute the loss between actual and predicted values
-    labels = torch.tensor([x[0] for x in labels]).to(device)
-    loss = cross_entropy(preds, labels)
-
-    # add on to the total loss
-    total_loss = total_loss + loss.item()
-
-    # backward pass to calculate the gradients
-    loss.backward()
-
-    # clip the the gradients to 1.0. It helps in preventing the exploding gradient problem
-    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
-
-    # update parameters
-    optimizer.step()
-
-    # model predictions are stored on GPU. So, push it to CPU
-    preds=preds.detach().cpu().numpy()
-
-    # append the model predictions
-    total_preds.append(preds)
-
-  # compute the training loss of the epoch
-  avg_loss = total_loss / len(train_dataloader)
-  
-  # predictions are in the form of (no. of batches, size of batch, no. of classes).
-  # reshape the predictions in form of (number of samples, no. of classes)
-  total_preds  = np.concatenate(total_preds, axis=0)
-
-  #returns the loss and predictions
-  return avg_loss, total_preds
-
-def evaluate():
-  
-  print("\nEvaluating...")
-  
-  # deactivate dropout layers
-  model.eval()
-
-  total_loss, total_accuracy = 0, 0
-  
-  # empty list to save the model predictions
-  total_preds = []
-
-  # iterate over batches
-  for step,batch in enumerate(train_dataloader):
-    
-    # Progress update every 50 batches.
-    if step % 50 == 0 and not step == 0:
-      
-      # Calculate elapsed time in minutes.
-            
-      # Report progress.
-      print('  Batch {:>5,}  of  {:>5,}.'.format(step, len(train_dataloader)))
-
-    # push the batch to gpu
-    batch = [t.to(device) for t in batch]
-
-    sent_id, mask, labels = batch
-
-    # deactivate autograd
-    with torch.no_grad():
-      
-      # model predictions
-      preds = model(sent_id, mask)
-
-      # compute the validation loss between actual and predicted values
-      labels = torch.tensor([x[0] for x in labels]).to(device)
-      loss = cross_entropy(preds,labels)
-
-      total_loss = total_loss + loss.item()
-
-      preds = preds.detach().cpu().numpy()
-
-      total_preds.append(preds)
-
-  # compute the validation loss of the epoch
-  avg_loss = total_loss / len(train_dataloader) 
-
-  # reshape the predictions in form of (number of samples, no. of classes)
-  total_preds  = np.concatenate(total_preds, axis=0)
-
-  return avg_loss, total_preds
-
-# avg_loss, total_preds = train()
-# set initial loss to infinite
-best_valid_loss = float('inf')
-
-# empty lists to store training and validation loss of each epoch
-train_losses=[]
-valid_losses=[]
-
-print("Started training!")
-#for each epoch
-for epoch in range(epochs):
-     
-    print('\n Epoch {:} / {:}'.format(epoch + 1, epochs))
-    
-    #train model
-    train_loss, _ = train()
-    
-    #evaluate model
-    valid_loss, _ = evaluate()
-    
-    #save the best model
-    if valid_loss < best_valid_loss:
-        best_valid_loss = valid_loss
-        torch.save(model.state_dict(), 'saved_weights.pt')
-    
-    # append training and validation loss
-    train_losses.append(train_loss)
-    valid_losses.append(valid_loss)
-    
-    print(f'\nTraining Loss: {train_loss:.3f}')
-    print(f'Validation Loss: {valid_loss:.3f}')
-
-print("Finished !!!")
\ No newline at end of file
+model_path = "bert-base-uncased-pretrained"
+model.save_pretrained(model_path)
+tokenizer.save_pretrained(model_path)
\ No newline at end of file
diff --git a/model.py b/model.py
deleted file mode 100644
index 9d218ce..0000000
--- a/model.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import torch.nn as nn
-
-class BERT_Arch(nn.Module):
-
-    def __init__(self, bert):
-      
-      super(BERT_Arch, self).__init__()
-
-      self.bert = bert 
-      
-      # dropout layer
-      self.dropout = nn.Dropout(0.1)
-      
-      # relu activation function
-      self.relu =  nn.ReLU()
-
-      # dense layer 1
-      self.fc1 = nn.Linear(2,512)
-      
-      # dense layer 2 (Output layer)
-      self.fc2 = nn.Linear(512,2)
-
-      #softmax activation function
-      self.softmax = nn.LogSoftmax(dim=1)
-
-    #define the forward pass
-    def forward(self, sent_id, mask):
-
-      #pass the inputs to the model  
-      senence_classifier_output = self.bert(sent_id, attention_mask=mask)
-      x = senence_classifier_output.logits.float()
-      x = self.fc1(x)
-
-      x = self.relu(x)
-
-      x = self.dropout(x)
-
-      # output layer
-      x = self.fc2(x)
-      
-      # apply softmax activation
-      x = self.softmax(x)
-
-      return x
\ No newline at end of file