Add .gitignore to ignore __pycache__ and other unnecessary files

.gitignore

@@ -0,0 +1,14 @@
+__pycache__/
+*.pyc
+*.pyo
+*.pyd
+.Python
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+.vscode/
+*.log
+*.sqlite3
+.DS_Store

src/models/advanced_architecture.py

@@ -2,22 +2,23 @@
 import torch.nn as nn
 import torch.nn.functional as F
 from typing import List, Tuple
+from torch.utils.data import DataLoader
 
 class InceptionModule(nn.Module):
     def __init__(self, in_channels: int, out_channels: List[int]):
         super(InceptionModule, self).__init__()
-        self.branch1x1 = nn.Conv2d(in_channels, out_channels[0], kernel_size=1)
+        self.branch1x1 = nn.Conv1d(in_channels, out_channels[0], kernel_size=1)
         self.branch3x3 = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels[1], kernel_size=1),
-            nn.Conv2d(out_channels[1], out_channels[2], kernel_size=3, padding=1)
+            nn.Conv1d(in_channels, out_channels[1], kernel_size=1),
+            nn.Conv1d(out_channels[1], out_channels[2], kernel_size=3, padding=1)
         )
         self.branch5x5 = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels[3], kernel_size=1),
-            nn.Conv2d(out_channels[3], out_channels[4], kernel_size=5, padding=2)
+            nn.Conv1d(in_channels, out_channels[3], kernel_size=1),
+            nn.Conv1d(out_channels[3], out_channels[4], kernel_size=5, padding=2)
         )
         self.branch_pool = nn.Sequential(
-            nn.MaxPool2d(kernel_size=3, stride=1, padding=1),
-            nn.Conv2d(in_channels, out_channels[5], kernel_size=1)
+            nn.MaxPool1d(kernel_size=3, stride=1, padding=1),
+            nn.Conv1d(in_channels, out_channels[5], kernel_size=1)
         )
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -37,12 +38,30 @@ def __init__(self, n_qubits: int, n_layers: int):
         self.entanglement = nn.Parameter(torch.randn(n_layers, n_qubits - 1))
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
+        # Ensure input tensor has the correct shape
+        if len(x.shape) == 2:
+            x = x.unsqueeze(0)  # Add batch dimension if not present
+        batch_size, seq_len, input_dim = x.shape
+
+        # Adjust input dimension to match n_qubits
+        if input_dim < self.n_qubits:
+            x = F.pad(x, (0, self.n_qubits - input_dim), "constant", 0)
+        elif input_dim > self.n_qubits:
+            x = x[:, :, :self.n_qubits]
+
         # Simplified quantum circuit simulation
         for layer in range(self.n_layers):
-            x = torch.sin(x + self.rotation[layer])
-            x = F.pad(x, (0, 1))
-            x = torch.roll(x, 1, dims=-1)
-            x = x[:, :-1] * torch.sin(self.entanglement[layer])
+            rotation = self.rotation[layer].unsqueeze(0).expand(batch_size, seq_len, -1, -1)
+            entanglement = self.entanglement[layer].unsqueeze(0).expand(batch_size, seq_len, -1)
+
+            # Apply rotation
+            x = torch.sin(x.unsqueeze(-1) + rotation).sum(dim=-1)
+
+            # Apply entanglement
+            x_shifted = torch.roll(x, 1, dims=-1)
+            x = x[:, :, :-1] * torch.sin(entanglement) + x_shifted[:, :, :-1] * torch.cos(entanglement)
+            x = F.pad(x, (0, 1), "constant", 0)  # Pad the last qubit
+
         return x
 
 class GraphNeuralNetwork(nn.Module):
@@ -92,8 +111,9 @@ def forward(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, m
         return output
 
 class AdvancedNeuroCoder(nn.Module):
-    def __init__(self, vocab_size: int, d_model: int = 768, n_layers: int = 12, n_heads: int = 12):
+    def __init__(self, vocab_size: int, d_model: int = 768, n_layers: int = 12, n_heads: int = 12, num_tasks: int = 3):
         super(AdvancedNeuroCoder, self).__init__()
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         self.embedding = nn.Embedding(vocab_size, d_model)
         self.inception = InceptionModule(d_model, [64, 96, 128, 16, 32, 32])
         self.quantum_layer = QuantumLayer(n_qubits=d_model, n_layers=2)
@@ -107,19 +127,73 @@ def __init__(self, vocab_size: int, d_model: int = 768, n_layers: int = 12, n_he
         self.layer_norm1 = nn.LayerNorm(d_model)
         self.layer_norm2 = nn.LayerNorm(d_model)
         self.output = nn.Linear(d_model, vocab_size)
+        self.task_classifier = nn.Linear(d_model, num_tasks)
+        self.criterion = nn.CrossEntropyLoss()
+        self.to(self.device)
 
-    def forward(self, x: torch.Tensor, adj_matrix: torch.Tensor = None) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, attention_mask: torch.Tensor = None, adj_matrix: torch.Tensor = None) -> Tuple[torch.Tensor, torch.Tensor]:
+        x = x.to(self.device)
         x = self.embedding(x)
-        x = x.unsqueeze(2).expand(-1, -1, x.size(1), -1)  # Expand for Inception
-        x = self.inception(x).squeeze(2)
+
+        # Ensure x has the correct shape for inception layer
+        batch_size, seq_len, d_model = x.shape
+        x = x.transpose(1, 2)  # Change shape to (batch, d_model, seq_len)
+
+        x = self.inception(x)
+        x = x.transpose(1, 2)  # Change shape back to (batch, seq_len, d_model)
+
+        # Handle potential shape mismatch in quantum layer
         x = self.quantum_layer(x)
+
         if adj_matrix is not None:
+            adj_matrix = adj_matrix.to(self.device)
             x = self.gnn(x, adj_matrix)
-        attn_output = self.attention(x, x, x)
+
+        # Apply attention mask if provided
+        if attention_mask is not None:
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = (1.0 - attention_mask) * -10000.0
+
+        attn_output = self.attention(x, x, x, mask=attention_mask)
         x = self.layer_norm1(x + attn_output)
         ff_output = self.feed_forward(x)
         x = self.layer_norm2(x + ff_output)
-        return self.output(x)
+        token_output = self.output(x)
+        task_output = self.task_classifier(x.mean(dim=1))  # Global average pooling
+
+        return token_output, task_output
+
+    def eval_loss(self, val_loader: DataLoader) -> float:
+        self.eval()
+        total_loss = 0
+        num_batches = 0
+        with torch.no_grad():
+            for batch in val_loader:
+                try:
+                    input_ids = batch['input_ids'].to(self.device)
+                    attention_mask = batch['attention_mask'].to(self.device)
+                    labels = batch['labels'].to(self.device)
+                    task_labels = batch['task_labels'].to(self.device)
+
+                    token_outputs, task_outputs = self(input_ids, attention_mask)
+
+                    # Ensure token_outputs and labels have the same shape
+                    if token_outputs.shape[1] != labels.shape[1]:
+                        min_len = min(token_outputs.shape[1], labels.shape[1])
+                        token_outputs = token_outputs[:, :min_len, :]
+                        labels = labels[:, :min_len]
+
+                    token_loss = self.criterion(token_outputs.contiguous().view(-1, token_outputs.size(-1)), labels.contiguous().view(-1))
+                    task_loss = self.criterion(task_outputs, task_labels)
+                    loss = token_loss + task_loss
+                    total_loss += loss.item()
+                    num_batches += 1
+                except RuntimeError as e:
+                    print(f"Error during evaluation: {e}")
+                    print(f"Input shape: {input_ids.shape}, Attention mask shape: {attention_mask.shape}")
+                    print(f"Labels shape: {labels.shape}, Task labels shape: {task_labels.shape}")
+                    continue
+        return total_loss / num_batches if num_batches > 0 else float('inf')
 
 # Example usage
 if __name__ == "__main__":

src/models/model_training.py

@@ -9,70 +9,145 @@
 from stable_baselines3.common.vec_env import DummyVecEnv
 from src.models.advanced_architecture import AdvancedNeuroCoder
 from bayes_opt import BayesianOptimization
+from torch.nn.utils.rnn import pad_sequence
 
 from src.models.advanced_architecture import AdvancedNeuroCoder
 
 # AdvancedNeuroCoder is now imported and will be used instead of the previous NeuroCoder class
 
 def load_datasets():
-    # TODO: Implement loading of large datasets
-    # Include code snippets, bug reports, and project documentation
-    # Ensure diverse programming languages and coding styles
-    pass
+    # Mock implementation for testing purposes
+    sequence_length = 100
+    task_to_label = {'code_generation': 0, 'bug_fixing': 1}
+    train_data = [
+        {'input_ids': torch.randint(0, 10000, (sequence_length,)), 'attention_mask': torch.ones(sequence_length), 'labels': torch.randint(0, 10000, (sequence_length,)), 'task': 'code_generation', 'task_labels': torch.tensor(task_to_label['code_generation'])},
+        {'input_ids': torch.randint(0, 10000, (sequence_length,)), 'attention_mask': torch.ones(sequence_length), 'labels': torch.randint(0, 10000, (sequence_length,)), 'task': 'bug_fixing', 'task_labels': torch.tensor(task_to_label['bug_fixing'])}
+    ]
+    val_data = [
+        {'input_ids': torch.randint(0, 10000, (sequence_length,)), 'attention_mask': torch.ones(sequence_length), 'labels': torch.randint(0, 10000, (sequence_length,)), 'task': 'code_generation', 'task_labels': torch.tensor(task_to_label['code_generation'])},
+        {'input_ids': torch.randint(0, 10000, (sequence_length,)), 'attention_mask': torch.ones(sequence_length), 'labels': torch.randint(0, 10000, (sequence_length,)), 'task': 'bug_fixing', 'task_labels': torch.tensor(task_to_label['bug_fixing'])}
+    ]
+    return train_data, val_data
 
 def generate_synthetic_data():
-    # TODO: Implement synthetic data generation
-    # Cover edge cases and uncommon scenarios
-    pass
+    # Mock implementation for testing purposes
+    sequence_length = 100
+    task_to_label = {'edge_case': 2, 'uncommon_scenario': 3}  # Continuing from previous task labels
+    return [
+        {'input_ids': torch.randint(0, 10000, (sequence_length,)), 'attention_mask': torch.ones(sequence_length), 'labels': torch.randint(0, 10000, (sequence_length,)), 'task': 'edge_case', 'task_labels': torch.tensor(task_to_label['edge_case'])},
+        {'input_ids': torch.randint(0, 10000, (sequence_length,)), 'attention_mask': torch.ones(sequence_length), 'labels': torch.randint(0, 10000, (sequence_length,)), 'task': 'uncommon_scenario', 'task_labels': torch.tensor(task_to_label['uncommon_scenario'])}
+    ]
 
 def train_model(model: AdvancedNeuroCoder, train_loader: DataLoader, val_loader: DataLoader, config: Dict[str, Any]):
     optimizer = AdamW(model.parameters(), lr=config['learning_rate'], weight_decay=config['weight_decay'])
     scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=config['warmup_steps'], num_training_steps=config['total_steps'])
-    criterion = nn.CrossEntropyLoss()
-    ppo = PPO(model, config['ppo_clip_param'], config['ppo_epochs'], config['ppo_batch_size'])
+    token_criterion = nn.CrossEntropyLoss(ignore_index=-100)  # Use -100 as padding index
+    task_criterion = nn.CrossEntropyLoss()
 
     for epoch in range(config['num_epochs']):
         model.train()
         total_loss = 0
-        for batch in train_loader:
-            input_ids = batch['input_ids']
-            attention_mask = batch['attention_mask']
-            labels = batch['labels']
-            task = batch['task']
-
-            # Generate actions (predictions) and calculate log probabilities
-            actions, log_probs = model(input_ids=input_ids, attention_mask=attention_mask, task=task)
-
-            # Calculate rewards (e.g., based on accuracy or other metrics)
-            rewards = calculate_rewards(actions, labels)
-
-            # Update the model using PPO
-            ppo_loss = ppo.update(input_ids, attention_mask, task, actions, log_probs, rewards)
-            total_loss += ppo_loss.item()
-
-            optimizer.step()
-            scheduler.step()
+        for batch_idx, batch in enumerate(train_loader):
+            optimizer.zero_grad()
+            input_ids = batch['input_ids'].to(model.device)
+            attention_mask = batch['attention_mask'].to(model.device)
+            labels = batch['labels'].to(model.device)
+            task_labels = batch['task_labels'].to(model.device)
+
+            try:
+                # Ensure input tensors have the correct shape
+                input_ids = input_ids.unsqueeze(0) if input_ids.dim() == 1 else input_ids
+                attention_mask = attention_mask.unsqueeze(0) if attention_mask.dim() == 1 else attention_mask
+                labels = labels.unsqueeze(0) if labels.dim() == 1 else labels
+                task_labels = task_labels.unsqueeze(0) if task_labels.dim() == 1 else task_labels
+
+                # Log input shapes for debugging
+                if batch_idx % 100 == 0:
+                    print(f"Batch {batch_idx}: Input shape: {input_ids.shape}, Attention mask shape: {attention_mask.shape}")
+
+                token_output, task_output = model(input_ids, attention_mask)
+
+                # Ensure token_output and labels have the same shape
+                if token_output.shape[1] != labels.shape[1]:
+                    min_len = min(token_output.shape[1], labels.shape[1])
+                    token_output = token_output[:, :min_len, :]
+                    labels = labels[:, :min_len]
+
+                # Mask out padding tokens
+                mask = (labels != -100).float()
+                token_loss = token_criterion(token_output.contiguous().view(-1, token_output.size(-1)), labels.contiguous().view(-1))
+                token_loss = (token_loss * mask.view(-1)).sum() / mask.sum()
+
+                task_loss = task_criterion(task_output, task_labels.squeeze())
+                loss = token_loss + task_loss
+
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(model.parameters(), config['max_grad_norm'])
+                optimizer.step()
+                scheduler.step()
+
+                total_loss += loss.item()
+
+                # Log detailed information for debugging
+                if batch_idx % 100 == 0:
+                    print(f"Batch {batch_idx}: Token Loss: {token_loss.item():.4f}, Task Loss: {task_loss.item():.4f}")
+                    print(f"Token Output Shape: {token_output.shape}, Labels Shape: {labels.shape}")
+                    print(f"Task Output Shape: {task_output.shape}, Task Labels Shape: {task_labels.shape}")
+
+            except RuntimeError as e:
+                print(f"Error during training (Batch {batch_idx}): {e}")
+                print(f"Input shape: {input_ids.shape}, Attention mask shape: {attention_mask.shape}")
+                print(f"Labels shape: {labels.shape}, Task labels shape: {task_labels.shape}")
+                continue
 
         avg_train_loss = total_loss / len(train_loader)
 
         # Validation
         model.eval()
         total_val_loss = 0
         with torch.no_grad():
-            for batch in val_loader:
-                input_ids = batch['input_ids']
-                attention_mask = batch['attention_mask']
-                labels = batch['labels']
-                task = batch['task']
-
-                outputs, _ = model(input_ids=input_ids, attention_mask=attention_mask, task=task)
-                loss = criterion(outputs.view(-1, outputs.size(-1)), labels.view(-1))
-                total_val_loss += loss.item()
+            for batch_idx, batch in enumerate(val_loader):
+                input_ids = batch['input_ids'].to(model.device)
+                attention_mask = batch['attention_mask'].to(model.device)
+                labels = batch['labels'].to(model.device)
+                task_labels = batch['task_labels'].to(model.device)
+
+                try:
+                    # Ensure input tensors have the correct shape
+                    input_ids = input_ids.unsqueeze(0) if input_ids.dim() == 1 else input_ids
+                    attention_mask = attention_mask.unsqueeze(0) if attention_mask.dim() == 1 else attention_mask
+                    labels = labels.unsqueeze(0) if labels.dim() == 1 else labels
+                    task_labels = task_labels.unsqueeze(0) if task_labels.dim() == 1 else task_labels
+
+                    token_output, task_output = model(input_ids, attention_mask)
+
+                    # Ensure token_output and labels have the same shape
+                    if token_output.shape[1] != labels.shape[1]:
+                        min_len = min(token_output.shape[1], labels.shape[1])
+                        token_output = token_output[:, :min_len, :]
+                        labels = labels[:, :min_len]
+
+                    # Mask out padding tokens
+                    mask = (labels != -100).float()
+                    token_loss = token_criterion(token_output.contiguous().view(-1, token_output.size(-1)), labels.contiguous().view(-1))
+                    token_loss = (token_loss * mask.view(-1)).sum() / mask.sum()
+
+                    task_loss = task_criterion(task_output, task_labels.squeeze())
+                    loss = token_loss + task_loss
+
+                    total_val_loss += loss.item()
+                except RuntimeError as e:
+                    print(f"Error during validation (Batch {batch_idx}): {e}")
+                    print(f"Input shape: {input_ids.shape}, Attention mask shape: {attention_mask.shape}")
+                    print(f"Labels shape: {labels.shape}, Task labels shape: {task_labels.shape}")
+                    continue
 
         avg_val_loss = total_val_loss / len(val_loader)
 
         print(f"Epoch {epoch+1}/{config['num_epochs']}, Train Loss: {avg_train_loss:.4f}, Validation Loss: {avg_val_loss:.4f}")
 
+    return model
+
 def calculate_rewards(actions, labels):
     # Implement reward calculation based on the task and performance
     # This is a placeholder implementation
@@ -143,7 +218,8 @@ def continuous_learning(model: AdvancedNeuroCoder, new_data: List[Dict[str, torc
     model.old_params = [param.clone().detach() for param in model.parameters()]
 
 if __name__ == "__main__":
-    model = AdvancedNeuroCoder(vocab_size=10000)  # Adjust vocab_size as needed
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = AdvancedNeuroCoder(vocab_size=10000).to(device)  # Adjust vocab_size as needed
     train_data, val_data = load_datasets()
     synthetic_data = generate_synthetic_data()
 
@@ -159,7 +235,8 @@ def continuous_learning(model: AdvancedNeuroCoder, new_data: List[Dict[str, torc
         'warmup_steps': 1000,
         'total_steps': 100000,
         'num_epochs': 10,
-        'max_grad_norm': 1.0
+        'max_grad_norm': 1.0,
+        'device': device
     }
 
     # Hyperparameter optimization

tests/test_integration.py

@@ -14,13 +14,23 @@ def model():
 
 def test_api_generate_code(client, model):
     input_data = {
-        "input_ids": [1, 2, 3, 4],
-        "attention_mask": [1, 1, 1, 1],
+        "input_ids": [[1, 2, 3, 4]],  # Add batch dimension
+        "attention_mask": [[1, 1, 1, 1]],  # Add batch dimension
         "task": "generate"
     }
     response = client.post("/generate-code", json=input_data)
     assert response.status_code == 200
-    assert "output" in response.json()
+    assert "token_output" in response.json()
+    assert "task_output" in response.json()
+
+    token_output = response.json()["token_output"]
+    task_output = response.json()["task_output"]
+
+    assert isinstance(token_output, list)
+    assert isinstance(task_output, list)
+    assert len(token_output) == 1  # Batch size
+    assert len(token_output[0]) == 4  # Sequence length
+    assert len(task_output) == 1  # Batch size
 
 def test_api_feedback(client):
     feedback_data = {