Apply code formatting and style improvements

- Applied black formatting to test files for consistent style
- Improved import formatting in test files (multi-line imports)
- Updated .gitignore and README.md
- Applied formatting to core modules and experiments
- All changes maintain functionality while improving code readability

Author: bangyen

.gitignore

@@ -76,4 +76,5 @@ Desktop.ini
 
 # Project-specific
 outputs/
-mlruns/emergent/
+mlruns/
+emergent/

README.md

@@ -59,6 +59,8 @@ The framework achieves state-of-the-art performance with advanced training techn
 
 ## Quick Start
 
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/bangyen/zsharp/blob/main/zsharp_demo.ipynb)
+
 ### Installation
 
 ```bash

emergent_demo.ipynb

@@ -94,9 +94,9 @@
    "source": [
     "# Configuration for our referential game\n",
     "config = CommunicationConfig(\n",
-    "    vocabulary_size=10,      # Size of the message vocabulary\n",
-    "    message_length=1,        # Length of messages (1 token for simplicity)\n",
-    "    hidden_size=64,          # Hidden layer size\n",
+    "    vocabulary_size=10,  # Size of the message vocabulary\n",
+    "    message_length=1,  # Length of messages (1 token for simplicity)\n",
+    "    hidden_size=64,  # Hidden layer size\n",
     ")\n",
     "\n",
     "# Learning parameters\n",
@@ -151,58 +151,68 @@
     "# Simple training loop for referential game\n",
     "from src.langlab.data.world import encode_object\n",
     "\n",
+    "\n",
     "def train_step(speaker, listener, scene_objects, target_idx, speaker_opt, listener_opt):\n",
     "    \"\"\"Single training step for the referential game.\"\"\"\n",
     "    speaker.train()\n",
     "    listener.train()\n",
-    "    \n",
+    "\n",
     "    # Encode the scene objects\n",
     "    scene_tensor = torch.stack([encode_object(obj) for obj in scene_objects]).to(device)\n",
-    "    target_tensor = torch.tensor([target_idx], dtype=torch.long).to(device)  # Add batch dimension\n",
-    "    \n",
+    "    target_tensor = torch.tensor([target_idx], dtype=torch.long).to(\n",
+    "        device\n",
+    "    )  # Add batch dimension\n",
+    "\n",
     "    # Speaker generates message about target object\n",
-    "    target_object = scene_tensor[target_idx:target_idx+1]\n",
-    "    message_logits, message_tokens, gesture_logits, gesture_tokens = speaker(target_object)\n",
-    "    \n",
+    "    target_object = scene_tensor[target_idx : target_idx + 1]\n",
+    "    message_logits, message_tokens, gesture_logits, gesture_tokens = speaker(\n",
+    "        target_object\n",
+    "    )\n",
+    "\n",
     "    # Add batch dimension to scene_tensor for Listener (batch_size=1, num_candidates=3, object_dim)\n",
     "    scene_tensor_batched = scene_tensor.unsqueeze(0)  # Shape: (1, 3, object_dim)\n",
-    "    \n",
+    "\n",
     "    # Listener tries to identify target from message and scene\n",
     "    listener_logits = listener(message_tokens, scene_tensor_batched)\n",
-    "    \n",
+    "\n",
     "    # Compute listener loss (cross-entropy)\n",
     "    listener_loss = torch.nn.functional.cross_entropy(listener_logits, target_tensor)\n",
-    "    \n",
+    "\n",
     "    # Compute speaker loss (REINFORCE-style)\n",
     "    # Get the predicted target from listener\n",
     "    predicted_target = torch.argmax(listener_logits, dim=1)\n",
     "    # Reward is 1 if correct, 0 if incorrect\n",
     "    reward = (predicted_target == target_tensor).float()\n",
-    "    \n",
+    "\n",
     "    # Speaker loss: encourage generating messages that lead to correct predictions\n",
     "    # Use the message logits to compute log probabilities\n",
     "    message_probs = torch.softmax(message_logits, dim=-1)\n",
-    "    message_log_probs = torch.log(message_probs + 1e-8)  # Add small epsilon for numerical stability\n",
-    "    \n",
+    "    message_log_probs = torch.log(\n",
+    "        message_probs + 1e-8\n",
+    "    )  # Add small epsilon for numerical stability\n",
+    "\n",
     "    # Get the log probability of the generated message\n",
-    "    speaker_log_prob = message_log_probs.gather(2, message_tokens.unsqueeze(-1)).squeeze(-1)\n",
+    "    speaker_log_prob = message_log_probs.gather(\n",
+    "        2, message_tokens.unsqueeze(-1)\n",
+    "    ).squeeze(-1)\n",
     "    speaker_log_prob = speaker_log_prob.sum(dim=1)  # Sum over message length\n",
-    "    \n",
+    "\n",
     "    # REINFORCE loss: -log_prob * reward\n",
     "    speaker_loss = -(speaker_log_prob * reward).mean()\n",
-    "    \n",
+    "\n",
     "    # Total loss\n",
     "    total_loss = listener_loss + 0.1 * speaker_loss  # Weight speaker loss less\n",
-    "    \n",
+    "\n",
     "    # Backpropagation\n",
     "    speaker_opt.zero_grad()\n",
     "    listener_opt.zero_grad()\n",
     "    total_loss.backward()\n",
     "    speaker_opt.step()\n",
     "    listener_opt.step()\n",
-    "    \n",
+    "\n",
     "    return total_loss.item()\n",
     "\n",
+    "\n",
     "# Train for a few steps\n",
     "print(\"Training agents...\")\n",
     "losses = []\n",
@@ -211,22 +221,34 @@
     "for step in tqdm(range(100), desc=\"Training\"):\n",
     "    # Sample a new scene each step\n",
     "    scene_objects, target_idx = sample_scene(k=3, seed=step)\n",
-    "    \n",
+    "\n",
     "    # Training step\n",
-    "    loss = train_step(speaker, listener, scene_objects, target_idx, \n",
-    "                     speaker_optimizer, listener_optimizer)\n",
+    "    loss = train_step(\n",
+    "        speaker,\n",
+    "        listener,\n",
+    "        scene_objects,\n",
+    "        target_idx,\n",
+    "        speaker_optimizer,\n",
+    "        listener_optimizer,\n",
+    "    )\n",
     "    losses.append(loss)\n",
-    "    \n",
+    "\n",
     "    # Evaluate accuracy every 10 steps\n",
     "    if step % 10 == 0:\n",
     "        speaker.eval()\n",
     "        listener.eval()\n",
     "        with torch.no_grad():\n",
-    "            scene_tensor = torch.stack([encode_object(obj) for obj in scene_objects]).to(device)\n",
-    "            target_object = scene_tensor[target_idx:target_idx+1]\n",
-    "            message_logits, message_tokens, gesture_logits, gesture_tokens = speaker(target_object)\n",
+    "            scene_tensor = torch.stack(\n",
+    "                [encode_object(obj) for obj in scene_objects]\n",
+    "            ).to(device)\n",
+    "            target_object = scene_tensor[target_idx : target_idx + 1]\n",
+    "            message_logits, message_tokens, gesture_logits, gesture_tokens = speaker(\n",
+    "                target_object\n",
+    "            )\n",
     "            # Add batch dimension to scene_tensor for Listener\n",
-    "            scene_tensor_batched = scene_tensor.unsqueeze(0)  # Shape: (1, 3, object_dim)\n",
+    "            scene_tensor_batched = scene_tensor.unsqueeze(\n",
+    "                0\n",
+    "            )  # Shape: (1, 3, object_dim)\n",
     "            listener_logits = listener(message_tokens, scene_tensor_batched)\n",
     "            predicted = torch.argmax(listener_logits, dim=1)\n",
     "            accuracy = (predicted == target_idx).float().mean().item()\n",
@@ -247,31 +269,38 @@
     "fig, ax = plt.subplots(1, 1, figsize=(10, 6))\n",
     "\n",
     "# Plot raw loss\n",
-    "ax.plot(losses, alpha=0.3, color='lightblue', label='Raw Loss')\n",
+    "ax.plot(losses, alpha=0.3, color=\"lightblue\", label=\"Raw Loss\")\n",
     "\n",
     "# Calculate and plot smoothed loss\n",
     "window_size = 10\n",
     "if len(losses) >= window_size:\n",
     "    smoothed_losses = []\n",
     "    for i in range(len(losses)):\n",
     "        start_idx = max(0, i - window_size + 1)\n",
-    "        smoothed_loss = sum(losses[start_idx:i+1]) / (i - start_idx + 1)\n",
+    "        smoothed_loss = sum(losses[start_idx : i + 1]) / (i - start_idx + 1)\n",
     "        smoothed_losses.append(smoothed_loss)\n",
-    "    \n",
-    "    ax.plot(smoothed_losses, color='blue', linewidth=2, label=f'Smoothed Loss (window={window_size})')\n",
     "\n",
-    "ax.set_title('Training Loss')\n",
-    "ax.set_xlabel('Step')\n",
-    "ax.set_ylabel('Loss')\n",
+    "    ax.plot(\n",
+    "        smoothed_losses,\n",
+    "        color=\"blue\",\n",
+    "        linewidth=2,\n",
+    "        label=f\"Smoothed Loss (window={window_size})\",\n",
+    "    )\n",
+    "\n",
+    "ax.set_title(\"Training Loss\")\n",
+    "ax.set_xlabel(\"Step\")\n",
+    "ax.set_ylabel(\"Loss\")\n",
     "ax.grid(True)\n",
     "ax.legend()\n",
     "\n",
     "plt.tight_layout()\n",
     "plt.show()\n",
     "\n",
-    "print(f\"Random baseline accuracy: {1/3:.2%} (3 objects)\")\n",
+    "print(f\"Random baseline accuracy: {1 / 3:.2%} (3 objects)\")\n",
     "print(f\"Final accuracy: {accuracies[-1]:.2%}\")\n",
-    "print(f\"Improvement over random: {(accuracies[-1] - 1/3)*100:.1f} percentage points\")"
+    "print(\n",
+    "    f\"Improvement over random: {(accuracies[-1] - 1 / 3) * 100:.1f} percentage points\"\n",
+    ")"
    ]
   },
   {

src/langlab/core/agents.py

@@ -330,8 +330,9 @@ def forward(
         Returns:
             Tensor of shape (batch_size, num_candidates) with probabilities over candidates.
         """
-        batch_size, num_candidates = candidate_objects.size(0), candidate_objects.size(
-            1
+        batch_size, num_candidates = (
+            candidate_objects.size(0),
+            candidate_objects.size(1),
         )
 
         # One-hot encode message tokens
@@ -710,8 +711,9 @@ def forward(
         Returns:
             Tensor of shape (batch_size, num_candidates) with probabilities over candidates.
         """
-        batch_size, num_candidates = candidate_objects.size(0), candidate_objects.size(
-            1
+        batch_size, num_candidates = (
+            candidate_objects.size(0),
+            candidate_objects.size(1),
         )
 
         # Embed message tokens

src/langlab/core/channel.py

@@ -130,8 +130,7 @@ def send_multimodal(
             )
         if gesture_size != self.config.gesture_size:
             raise ValueError(
-                f"Expected gesture_size={self.config.gesture_size}, "
-                f"got {gesture_size}"
+                f"Expected gesture_size={self.config.gesture_size}, got {gesture_size}"
             )
 
         # Sample tokens

src/langlab/experiments/contact.py

@@ -283,7 +283,9 @@ def measure_intelligibility(self) -> None:
 
         # Create evaluation dataset
         eval_dataset = ReferentialGameDataset(
-            n_scenes=1000, k=self.config.k, seed=42  # Fixed size for evaluation
+            n_scenes=1000,
+            k=self.config.k,
+            seed=42,  # Fixed size for evaluation
         )
         eval_dataloader = DataLoader(
             eval_dataset, batch_size=self.config.batch_size, shuffle=False

src/langlab/experiments/ensemble_training.py

@@ -137,7 +137,7 @@ def train_ensemble(
 
     for i in range(n_models):
         seed = base_seed + i * 1000
-        logger.info(f"Training model {i+1}/{n_models} with seed {seed}")
+        logger.info(f"Training model {i + 1}/{n_models} with seed {seed}")
 
         # Train individual model
         train(
@@ -202,9 +202,9 @@ def train_ensemble(
             )
             speaker.load_state_dict(checkpoint["speaker_state_dict"])
             listener.load_state_dict(checkpoint["listener_state_dict"])
-            logger.info(f"Loaded checkpoint for model {i+1}")
+            logger.info(f"Loaded checkpoint for model {i + 1}")
         else:
-            logger.warning(f"No checkpoint found for model {i+1}")
+            logger.warning(f"No checkpoint found for model {i + 1}")
 
         speakers.append(speaker)
         listeners.append(listener)

src/langlab/training/train_grounded.py

@@ -100,7 +100,6 @@ def update(self, success: bool) -> bool:
             and self.level_successes / self.episodes_at_level >= self.success_threshold
             and self.current_level < len(self.curriculum_grids) - 1
         ):
-
             logger.info(f"Advancing to curriculum level {self.current_level + 1}")
             self.current_level += 1
             self.episodes_at_level = 0

tests/integration/test_ablate.py

@@ -174,7 +174,6 @@ def test_run_ablation_suite_experiment_id_format() -> None:
     ) as mock_zipf, patch(
         "torch.load"
     ) as mock_load:
-
         mock_eval.return_value = {
             "train": {"acc": 0.8},
             "iid": {"acc": 0.75},

tests/integration/test_cli.py

@@ -214,7 +214,8 @@ def test_heldout_parsing(self) -> None:
         """Test heldout pair parsing in train command."""
         runner = CliRunner()
         result = runner.invoke(
-            main, ["train", "--steps", "10", "--heldout", "red,circle"]  # Valid format
+            main,
+            ["train", "--steps", "10", "--heldout", "red,circle"],  # Valid format
         )
         # Should not crash on parsing
         assert result.exit_code == 0 or result.exit_code == 1  # May fail on training