Update train_model.py

train_model.py

@@ -11,39 +11,30 @@ from tensor_network import FourDimensionalTransformer  # Adjust the import path
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-# List of dataset identifiers
+# List of dataset identifiers for reasoning and knowledge
 dataset_ids = [
-    "prithivMLmods/Deepthink-Reasoning",
-    "ewok-core/ewok-core-1.0",
-    "MuskumPillerum/General-Knowledge",
-    "fblgit/tree-of-knowledge",
-    "CohereForAI/aya_dataset",
-    "AtlasUnified/Atlas-Reasoning",
-    "livebench/reasoning",
-    "SkunkworksAI/reasoning-0.01",
-    "KingNish/reasoning-base-20k",
-    "RLHFlow/HH-RLHF-Helpful-standard",
-    "yitingxie/rlhf-reward-datasets"
+    "race/all",    # For reasoning
+    "squad"        # For general knowledge
 ]
 
+# Update possible keys
+possible_text_keys = ['question', 'sentence', 'query']
+possible_context_keys = ['context', 'article', 'passage']
+possible_label_keys = ['answer', 'answers', 'options']
+
 # Initialize tokenizer
 tokenizer = AutoTokenizer.from_pretrained('bert-base-uncased')
 
-def tokenize_function(examples):
-    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]
+def tokenize_function_race(examples):
+    texts = [q + " " + p for q, p in zip(examples['question'], examples['article'])]
+    labels = examples['answer']
+    tokenized_inputs = tokenizer(texts, padding='max_length', truncation=True, max_length=48)
+    tokenized_inputs['labels'] = labels
+    return tokenized_inputs
 
-    texts = [str(t) for t in examples[text_key]]
+def tokenize_function_squad(examples):
+    texts = [q + " " + c for q, c in zip(examples['question'], examples['context'])]
+    labels = [ans['text'][0] if ans['text'] else '' for ans in examples['answers']]
     tokenized_inputs = tokenizer(texts, padding='max_length', truncation=True, max_length=48)
     tokenized_inputs['labels'] = labels
     return tokenized_inputs
@@ -52,47 +43,81 @@ def tokenize_function(examples):
 label_encoder = LabelEncoder()
 all_labels = []
 
-# Process each dataset individually
+# Process RACE dataset
+race_dataset = load_dataset('race', 'all')
 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}")
+for split in race_dataset.keys():
+    tokenized_race = race_dataset[split].map(
+        tokenize_function_race,
+        batched=True,
+        remove_columns=race_dataset[split].column_names,
+        load_from_cache_file=False,
+    )
+    tokenized_datasets.append({split: tokenized_race})
+    # Collect labels
+    all_labels.extend(tokenized_race['labels'])
+
+# Process SQuAD dataset
+squad_dataset = load_dataset('squad')
+for split in squad_dataset.keys():
+    tokenized_squad = squad_dataset[split].map(
+        tokenize_function_squad,
+        batched=True,
+        remove_columns=squad_dataset[split].column_names,
+        load_from_cache_file=False,
+    )
+    tokenized_datasets.append({split: tokenized_squad})
+    # Collect labels
+    all_labels.extend(tokenized_squad['labels'])
 
 # Fit label encoder
 label_encoder.fit(all_labels)
 num_classes = len(label_encoder.classes_)
 print(f"Number of unique labels: {num_classes}")
 
+# Limit the number of classes to top 10 frequent labels
 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
-            )
+    print("Number of classes exceeds 10. Reducing to top 10 classes.")
+    from collections import Counter
+    label_counter = Counter(all_labels)
+    top_10_labels = [label for label, _ in label_counter.most_common(10)]
+    print(f"Top 10 labels: {top_10_labels}")
+    label_mapping = {label: i for i, label in enumerate(top_10_labels)}
+    label_mapping['other'] = len(top_10_labels)
+    num_classes = len(top_10_labels) + 1
+else:
+    label_mapping = {label: i for i, label in enumerate(label_encoder.classes_)}
+
+# Update model with correct num_classes
+model = FourDimensionalTransformer(
+    num_layers=16,
+    embed_dim=7,
+    num_heads=1,
+    num_extra_tokens=16,
+    num_classes=num_classes
+).to(device)
+
+def map_labels(labels):
+    return [label_mapping.get(label, label_mapping['other']) for label in labels]
+
+# Process datasets
+for tokenized_dataset in tokenized_datasets:
+    for split in tokenized_dataset.keys():
+        tokenized_dataset[split] = tokenized_dataset[split].map(
+            lambda examples: {'labels': map_labels(examples['labels'])},
+            batched=True
+        )
+        tokenized_dataset[split] = tokenized_dataset[split].filter(
+            lambda example: example['labels'] < num_classes
+        )
+        tokenized_dataset[split].set_format(type='torch', columns=['input_ids', 'labels'])
 
 # Prepare DataLoaders
-def prepare_dataloader(dataset_splits, split_name, batch_size=2):
+def prepare_dataloader(tokenized_datasets, split_name, batch_size=4):
     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'])
+    for tokenized_dataset in tokenized_datasets:
+        if split_name in tokenized_dataset:
+            dataset_split = tokenized_dataset[split_name]
             dataloader = DataLoader(dataset_split, batch_size=batch_size, shuffle=True)
             dataloaders.append(dataloader)
     return dataloaders
@@ -100,20 +125,11 @@ def prepare_dataloader(dataset_splits, split_name, batch_size=2):
 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  # Using 10 classes as per your model
-).to(device)
-
 # Loss function and optimizer
 criterion = nn.CrossEntropyLoss()
 optimizer = optim.Adam(model.parameters(), lr=1e-4)
 
-def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
+def train(model, train_dataloaders, val_dataloaders, num_epochs=10): #change number of Epochs to your liking
     for epoch in range(num_epochs):
         model.train()
         total_loss = 0
@@ -124,11 +140,11 @@ def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
                 labels = batch['labels']
 
                 # Reshape input_ids and move to device
-                input_ids = input_ids[:, :48]  # Ensure length is 48
+                input_ids = input_ids[:, :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)
+                labels = labels.to(device).long()
 
                 optimizer.zero_grad()
                 outputs = model(input_ids)
@@ -152,9 +168,9 @@ def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
                     input_ids = batch['input_ids']
                     labels = batch['labels']
 
-                    input_ids = input_ids[:, :48]  # Ensure length is 48
+                    input_ids = input_ids[:, :48]
                     input_ids = input_ids.view(-1, 3, 4, 4).float().to(device)
-                    labels = labels.type(torch.long).to(device)
+                    labels = labels.to(device).long()
 
                     outputs = model(input_ids)
                     _, predicted = torch.max(outputs, 1)
@@ -166,4 +182,7 @@ def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
     torch.save(model.state_dict(), 'trained_model.pth')
 
 # Start training
-train(model, train_dataloaders, val_dataloaders)
+if train_dataloaders and val_dataloaders:
+    train(model, train_dataloaders, val_dataloaders)
+else:
+    print("No data loaders available for training. Please check the datasets and preprocessing steps.")