Update train_model.py

train_model.py

@@ -4,7 +4,12 @@ import torch.optim as optim
 from torch.utils.data import DataLoader
 from datasets import load_dataset
 from transformers import AutoTokenizer
-from tensor_network import FourDimensionalTransformer  # Adjust based on your model's location
+from sklearn.preprocessing import LabelEncoder
+
+# Import your model from tensor_network.py
+from tensor_network import FourDimensionalTransformer  # Adjust the import path as needed
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # List of dataset identifiers
 dataset_ids = [
@@ -21,55 +26,88 @@ dataset_ids = [
     "yitingxie/rlhf-reward-datasets"
 ]
 
-# Load datasets
-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
+tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
 
 def tokenize_function(examples):
-    """
-    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]
-
-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)
-
-# 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
+    possible_text_keys = ['text', 'content', 'question', 'passage', 'prompt', 'input']
+    possible_label_keys = ['label', 'answer', 'response', 'output', 'target']
+
+    text_key = next((k for k in possible_text_keys if k in examples), None)
+    if text_key is None:
+        text_key = list(examples.keys())[0]
+
+    label_key = next((k for k in possible_label_keys if k in examples), None)
+    if label_key is None:
+        labels = [0] * len(examples[text_key])  # Default label
+    else:
+        labels = examples[label_key]
+
+    texts = [str(t) for t in examples[text_key]]
+    tokenized_inputs = tokenizer(texts, padding='max_length', truncation=True, max_length=48)
+    tokenized_inputs['labels'] = labels
+    return tokenized_inputs
+
+# Initialize LabelEncoder
+label_encoder = LabelEncoder()
+all_labels = []
+
+# Process each dataset individually
+tokenized_datasets = []
+for dataset_id in dataset_ids:
+    try:
+        dataset = load_dataset(dataset_id)
+        tokenized_dataset = dataset.map(tokenize_function, batched=True)
+
+        # Collect labels for label encoding
+        for split in tokenized_dataset.keys():
+            if 'labels' in tokenized_dataset[split].features:
+                all_labels.extend(tokenized_dataset[split]['labels'])
+
+        tokenized_datasets.append(tokenized_dataset)
+    except Exception as e:
+        print(f"Could not process dataset {dataset_id}: {e}")
+
+# Fit label encoder
+label_encoder.fit(all_labels)
+num_classes = len(label_encoder.classes_)
+print(f"Number of unique labels: {num_classes}")
+
+if num_classes > 10:
+    print("Warning: Number of unique labels exceeds the number of classes. Adjusting the dataset or model is required.")
+    exit()
+
+# Transform labels in each dataset
+for dataset in tokenized_datasets:
+    for split in dataset.keys():
+        if 'labels' in dataset[split].features:
+            dataset[split] = dataset[split].map(
+                lambda examples: {'labels': label_encoder.transform(examples['labels'])},
+                batched=True
+            )
+
+# Prepare DataLoaders
+def prepare_dataloader(dataset_splits, split_name, batch_size=2):
+    dataloaders = []
+    for dataset in dataset_splits:
+        if split_name in dataset:
+            dataset_split = dataset[split_name]
+            dataset_split.set_format(type='torch', columns=['input_ids', 'labels'])
+            dataloader = DataLoader(dataset_split, batch_size=batch_size, shuffle=True)
+            dataloaders.append(dataloader)
+    return dataloaders
+
+train_dataloaders = prepare_dataloader(tokenized_datasets, 'train')
+val_dataloaders = prepare_dataloader(tokenized_datasets, 'validation')
+
+# Initialize the model
 model = FourDimensionalTransformer(
     num_layers=16,
     embed_dim=7,
     num_heads=1,
     num_extra_tokens=16,
-    num_classes=10  # Adjust based on your specific task
-)
+    num_classes=10  # Using 10 classes as per your model
+).to(device)
 
 # Loss function and optimizer
 criterion = nn.CrossEntropyLoss()
@@ -79,22 +117,29 @@ 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)
+        total_batches = 0
         for dataloader in train_dataloaders:
             for batch in dataloader:
                 input_ids = batch['input_ids']
-                attention_mask = batch['attention_mask']
-                labels = batch['label']
+                labels = batch['labels']
+
+                # Reshape input_ids and move to device
+                input_ids = input_ids[:, :48]  # Ensure length is 48
+                input_ids = input_ids.view(-1, 3, 4, 4).float().to(device)
+
+                # Convert labels to torch.long and move to device
+                labels = labels.type(torch.long).to(device)
 
                 optimizer.zero_grad()
-                outputs = model(input_ids, attention_mask)
+                outputs = model(input_ids)
                 loss = criterion(outputs, labels)
                 loss.backward()
                 optimizer.step()
+
                 total_loss += loss.item()
+                total_batches += 1
 
-        # Use the last dataloader's length to compute average loss
-        avg_loss = total_loss / len(dataloader)
+        avg_loss = total_loss / total_batches if total_batches > 0 else 0
         print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}')
 
         # Validation loop
@@ -105,16 +150,19 @@ def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
             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)
+                    labels = batch['labels']
+
+                    input_ids = input_ids[:, :48]  # Ensure length is 48
+                    input_ids = input_ids.view(-1, 3, 4, 4).float().to(device)
+                    labels = labels.type(torch.long).to(device)
+
+                    outputs = model(input_ids)
                     _, predicted = torch.max(outputs, 1)
                     total_correct += (predicted == labels).sum().item()
                     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')
 
 # Start training