Reduce sample size and increase epochs

app.py

@@ -3,6 +3,7 @@ import streamlit as st
 from datasets import load_dataset
 
 import numpy as np
+import os
 from sklearn.metrics import accuracy_score, precision_recall_fscore_support
 
 import torch
@@ -13,6 +14,7 @@ from torch.utils.data import DataLoader
 from transformers import AutoModelForTokenClassification, AutoTokenizer, DataCollatorForTokenClassification
 from transformers import DebertaV2Config, DebertaV2ForTokenClassification
 
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" 
 
 # print weights
 def print_trainable_parameters(model):
@@ -44,7 +46,7 @@ st.write("dimension", dimension)
 
 student_model_config = teacher_model.config
 student_model_config.num_attention_heads = 8
-student_model_config.num_hidden_layers = 6
+student_model_config.num_hidden_layers = 4
 student_model = DebertaV2ForTokenClassification.from_pretrained(
     "microsoft/mdeberta-v3-base",
     config=student_model_config)
@@ -59,7 +61,7 @@ if torch.cuda.is_available():
 # Load data. 
 raw_dataset = load_dataset("ai4privacy/pii-masking-400k", split='train')
 raw_dataset = raw_dataset.filter(lambda example: example["language"].startswith("en"))
-raw_dataset = raw_dataset.select(range(4000)) 
+raw_dataset = raw_dataset.select(range(2000)) 
 raw_dataset = raw_dataset.train_test_split(test_size=0.2)
 print(raw_dataset)
 print(raw_dataset.column_names)
@@ -175,7 +177,7 @@ def distillation_loss(student_logits, teacher_logits, true_labels, temperature,
 # hyperparameters
 batch_size = 32
 lr = 1e-4
-num_epochs = 30
+num_epochs = 300
 temperature = 2.0
 alpha = 0.5