Upload train_model.py

train_model.py

@@ -0,0 +1,103 @@
+import torch
+import torch.nn as nn
+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
+
+# List of dataset identifiers
+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",
+    "MBZUAI/ArabicMMLU"
+]
+
+# 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
+
+# Tokenize datasets
+def tokenize_function(examples):
+    return tokenizer(examples['text'], padding='max_length', truncation=True, max_length=128)
+
+tokenized_datasets = [dataset.map(tokenize_function, batched=True) for dataset in datasets]
+
+
+# Prepare DataLoader
+def prepare_dataloader(dataset, batch_size=32):
+    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]
+
+
+# Model setup
+model = FourDimensionalTransformer(
+    num_layers=16,
+    embed_dim=7,
+    num_heads=1,
+    num_extra_tokens=16,
+    num_classes=10  # Adjust based on your specific task
+)
+
+# 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
+
+# Training loop
+def train(model, train_dataloaders, val_dataloaders, num_epochs=10):
+    for epoch in range(num_epochs):
+        model.train()
+        total_loss = 0
+        for dataloader in train_dataloaders:
+            for batch in dataloader:
+                input_ids = batch['input_ids']
+                attention_mask = batch['attention_mask']
+                labels = batch['label']
+
+                optimizer.zero_grad()
+                outputs = model(input_ids, attention_mask)
+                loss = criterion(outputs, labels)
+                loss.backward()
+                optimizer.step()
+
+                total_loss += loss.item()
+
+        avg_loss = total_loss / len(dataloader)
+        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {avg_loss:.4f}')
+
+        # Validation
+        model.eval()
+        total_correct = 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)
+        print(f'Validation Accuracy: {accuracy:.4f}')
+
+    # Save the trained model
+    torch.save(model.state_dict(), 'trained_model.pth')
+
+
+# Train the model
+train(model, train_dataloaders, val_dataloaders)