Update app.py

app.py

@@ -1,7 +1,7 @@
 import streamlit as st
 import torch
 from torch.utils.data import Dataset, random_split
-from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, default_data_collator
+from transformers import AutoModelForSequenceClassification, AutoTokenizer, Trainer, TrainingArguments, DataCollatorForLanguageModeling
 
 # Generate sample clinical text and labels
 sample_data = [
@@ -58,7 +58,10 @@ class ClinicalDataset(Dataset):
 
 
 # Data Collator
-data_collator = default_data_collator
+data_collator = DataCollatorForLanguageModeling(
+    tokenizer=tokenizer, 
+    mlm_probability=0.15
+)
 
 seq_length = 128
 dataset = ClinicalDataset(texts=train_texts, labels=train_labels, tokenizer=tokenizer, max_seq_length=seq_length)
@@ -84,11 +87,7 @@ trainer = Trainer(
     args=training_args,
     train_dataset=train_dataset,
     eval_dataset=val_dataset,
-    data_collator=lambda data: {'input_ids': torch.stack([f['input_ids'] for f in data]),
-                            'attention_mask': torch.stack([f['attention_mask'] for f in data]),
-                             'labels': torch.stack([f['labels'] for f in data])}, 
-    pad_to_max_length=True
-
+    data_collator=data_collator, 
 )