Update train_model.py

train_model.py

@@ -27,21 +27,40 @@ datasets = [load_dataset(dataset_id) for dataset_id in dataset_ids]
 # Initialize tokenizer
 tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')  # Replace with your model's tokenizer
 
-# Tokenize datasets
 def tokenize_function(examples):
-    return tokenizer(examples['text'], padding='max_length', truncation=True, max_length=128)
-
+    """
+    Attempts to find a common text column (e.g., 'text', 'content', 'question', or 'passage').
+    If not found, it falls back to the first available key.
+    Ensures that all inputs are converted to strings.
+    """
+    possible_keys = ['text', 'content', 'question', 'passage']
+    key = None
+    for k in possible_keys:
+        if k in examples:
+            key = k
+            break
+    if key is None:
+        key = list(examples.keys())[0]  # fallback if none of the common keys exist
+    # Convert all items to string in case they're not
+    texts = [str(t) for t in examples[key]]
+    return tokenizer(texts, padding='max_length', truncation=True, max_length=128)
+
+# Apply tokenization to all datasets
 tokenized_datasets = [dataset.map(tokenize_function, batched=True) for dataset in datasets]
 
-
-# Prepare DataLoader
 def prepare_dataloader(dataset, batch_size=32):
+    """
+    Sets the format for the dataset to PyTorch and returns a DataLoader.
+    This function assumes the dataset contains 'input_ids', 'attention_mask', and 'label'.
+    If the label column is missing, you'll need to adjust this accordingly.
+    """
+    # You may need to adjust the columns if your datasets use a different label column name.
     dataset.set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])
     return DataLoader(dataset, batch_size=batch_size, shuffle=True)
 
-train_dataloaders = [prepare_dataloader(dataset['train']) for dataset in tokenized_datasets]
-val_dataloaders = [prepare_dataloader(dataset['validation']) for dataset in tokenized_datasets]
-
+# Only include splits that exist to avoid key errors
+train_dataloaders = [prepare_dataloader(ds['train']) for ds in tokenized_datasets if 'train' in ds]
+val_dataloaders = [prepare_dataloader(ds['validation']) for ds in tokenized_datasets if 'validation' in ds]
 
 # Model setup
 model = FourDimensionalTransformer(
@@ -54,13 +73,13 @@ model = FourDimensionalTransformer(
 
 # Loss function and optimizer
 criterion = nn.CrossEntropyLoss()
-optimizer = optim.Adam(model.parameters(), lr=1e-4)  # Using Adam optimizer with a learning rate of 1e-4
+optimizer = optim.Adam(model.parameters(), lr=1e-4)
 
-# Training loop
 def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
     for epoch in range(num_epochs):
         model.train()
         total_loss = 0
+        # Iterate over each training dataloader (from each dataset)
         for dataloader in train_dataloaders:
             for batch in dataloader:
                 input_ids = batch['input_ids']
@@ -72,32 +91,31 @@ def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
                 loss = criterion(outputs, labels)
                 loss.backward()
                 optimizer.step()
-
                 total_loss += loss.item()
 
+        # Use the last dataloader's length to compute average loss
         avg_loss = total_loss / len(dataloader)
         print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}')
 
-        # Validation
+        # Validation loop
         model.eval()
         total_correct = 0
+        total_samples = 0
         with torch.no_grad():
             for dataloader in val_dataloaders:
                 for batch in dataloader:
                     input_ids = batch['input_ids']
                     attention_mask = batch['attention_mask']
                     labels = batch['label']
-
                     outputs = model(input_ids, attention_mask)
                     _, predicted = torch.max(outputs, 1)
                     total_correct += (predicted == labels).sum().item()
-
-        accuracy = total_correct / len(dataloader.dataset)
+                    total_samples += labels.size(0)
+        accuracy = total_correct / total_samples if total_samples > 0 else 0
         print(f'Validation Accuracy: {accuracy:.4f}')
 
     # Save the trained model
     torch.save(model.state_dict(), 'trained_model.pth')
 
-
-# Train the model
+# Start training
 train(model, train_dataloaders, val_dataloaders)