Update app to set tensors to consistent size

app.py

@@ -28,10 +28,11 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Create PyTorch Dataset object
 class ClinicalDataset(Dataset):
-         def __init__(self, texts, labels, tokenizer):
+         def __init__(self, texts, labels, tokenizer, max_seq_length):
              self.texts = texts
              self.labels = labels
              self.tokenizer = tokenizer
+             self.max_seq_length = max_seq_length
      
          def __len__(self):
              return len(self.texts)
@@ -39,15 +40,28 @@ class ClinicalDataset(Dataset):
          def __getitem__(self, idx):
              text = self.texts[idx]
              label = self.labels[idx]
-             encoding = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
-             return {"input_ids": encoding["input_ids"].squeeze(), "attention_mask": encoding["attention_mask"].squeeze(), "labels": torch.tensor(label)}
+             
+             encoding = self.tokenizer(
+                text, 
+                return_tensors="pt", 
+                padding='max_length',  # Pad sequences to the maximum sequence length
+                truncation=True, 
+                max_length=self.max_seq_length
+             )
+             
+             return {
+                 "input_ids": encoding["input_ids"].squeeze(), 
+                 "attention_mask": encoding["attention_mask"].squeeze(), 
+                 "labels": torch.tensor(label)
+             }
      
 
 
 # Data Collator
 data_collator = default_data_collator
 
-dataset = ClinicalDataset(texts=train_texts, labels=train_labels, tokenizer=tokenizer)
+seq_length = 128
+dataset = ClinicalDataset(texts=train_texts, labels=train_labels, tokenizer=tokenizer, max_seq_length=seq_length)
 
 # Split dataset into training and validation sets
 train_size = int(0.8 * len(dataset))
@@ -66,13 +80,16 @@ training_args = TrainingArguments(
          logging_steps=10,)
 
 trainer = Trainer(
-      model=model,
-      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])}, )
+    model=model,
+    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
+
+)
 
 
 st.write("Training started...")