feat: integrate MAAS framework into MASLab with dual mode support and…

inference.py

@@ -130,7 +130,7 @@ def reserve_unprocessed(output_json, test_dataset):
         lock = threading.Lock()
         if args.sequential:
             for sample in test_dataset:
-                process_sample(args, general_config, sample, output_path)
+                process_sample(args, general_config, sample, output_path, lock)
         else:
             max_workers = model_api_config[args.model_name]["max_workers"]
             with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:

methods/__init__.py

@@ -1,59 +1,91 @@
-from .mas_base import MAS
-from .cot import CoT
-from .agentverse import AgentVerse_HumanEval, AgentVerse_MGSM, AgentVerse_Main
-from .llm_debate import LLM_Debate_Main
-from .dylan import DyLAN_HumanEval, DyLAN_MATH, DyLAN_MMLU, DyLAN_Main
-from .autogen import AutoGen_Main
-from .camel import CAMEL_Main
-from .evomac import EvoMAC_Main
-from .chatdev import ChatDev_SRDD
-from .macnet import MacNet_Main, MacNet_SRDD
-from .mad import MAD_Main
-from .mapcoder import MapCoder_HumanEval, MapCoder_MBPP
-from .self_consistency import SelfConsistency
-
-method2class = {
-    "vanilla": MAS,
-    "cot": CoT,
-    "agentverse_humaneval": AgentVerse_HumanEval,
-    "agentverse_mgsm": AgentVerse_MGSM,
-    "agentverse": AgentVerse_Main,
-    "llm_debate": LLM_Debate_Main,
-    "dylan_humaneval": DyLAN_HumanEval,
-    "dylan_math": DyLAN_MATH,
-    "dylan_mmlu": DyLAN_MMLU,
-    "dylan": DyLAN_Main,
-    "autogen": AutoGen_Main,
-    "camel": CAMEL_Main,
-    "evomac": EvoMAC_Main,
-    "chatdev_srdd": ChatDev_SRDD,
-    "macnet": MacNet_Main,
-    "macnet_srdd": MacNet_SRDD,
-    "mad": MAD_Main,
-    "mapcoder_humaneval": MapCoder_HumanEval,
-    "mapcoder_mbpp": MapCoder_MBPP,
-    "self_consistency": SelfConsistency_Main
+import importlib
+from types import ModuleType
+from typing import Dict, Tuple
+
+__all__ = ["get_method_class"]
+
+# ---------------------------------------------------------------------------
+# Lazy registry: mapping from *public* method name (argument to --method_name)
+# to a tuple (module_path, class_name).  The actual import happens only when
+# the user explicitly selects that method via get_method_class().  This avoids
+# importing optional heavy dependencies at MASLab startup.
+# ---------------------------------------------------------------------------
+_method_registry: Dict[str, Tuple[str, str]] = {
+    # Core methods that have minimal external deps
+    "vanilla": ("methods.mas_base", "MAS"),
+    "cot": ("methods.cot", "CoT"),
+
+    # AgentVerse variants
+    "agentverse": ("methods.agentverse.agentverse_main", "AgentVerse_Main"),
+    "agentverse_humaneval": ("methods.agentverse.agentverse_humaneval", "AgentVerse_HumanEval"),
+    "agentverse_mgsm": ("methods.agentverse.agentverse_mgsm", "AgentVerse_MGSM"),
+
+    # Maas integration (the focus of current work)
+    "maas": ("methods.maas.maas_main", "MAAS_Main"),
+
+    # The remaining methods are kept but lazily imported.  They may rely on
+    # optional third-party libraries which might not be installed in every
+    # environment.  Users can install the required dependencies when they need
+    # these methods.
+    "llm_debate": ("methods.llm_debate", "LLM_Debate_Main"),
+    "autogen": ("methods.autogen.autogen_main", "AutoGen_Main"),
+    "camel": ("methods.camel", "CAMEL_Main"),
+    "evomac": ("methods.evomac", "EvoMAC_Main"),
+    "chatdev_srdd": ("methods.chatdev.chatdev_srdd", "ChatDev_SRDD"),
+    "macnet": ("methods.macnet.macnet_main", "MacNet_Main"),
+    "macnet_srdd": ("methods.macnet.macnet_srdd", "MacNet_SRDD"),
+    "mad": ("methods.mad", "MAD_Main"),
+    "mapcoder_humaneval": ("methods.mapcoder.mapcoder_humaneval", "MapCoder_HumanEval"),
+    "mapcoder_mbpp": ("methods.mapcoder.mapcoder_mbpp", "MapCoder_MBPP"),
+    # DyLAN and SelfConsistency variants
+    "dylan": ("methods.dylan.dylan_main", "DyLAN_Main"),
+    "dylan_humaneval": ("methods.dylan.dylan_humaneval", "DyLAN_HumanEval"),
+    "dylan_math": ("methods.dylan.dylan_math", "DyLAN_MATH"),
+    "dylan_mmlu": ("methods.dylan.dylan_mmlu", "DyLAN_MMLU"),
+    "self_consistency": ("methods.self_consistency", "SelfConsistency"),
 }
 
-def get_method_class(method_name, dataset_name=None):
-
-    # lowercase the method name
+
+def _import_class(module_path: str, class_name: str):
+    """Import *class_name* from *module_path* and return the class object."""
+    module: ModuleType = importlib.import_module(module_path)
+    try:
+        return getattr(module, class_name)
+    except AttributeError as exc:
+        raise ImportError(f"Class '{class_name}' not found in module '{module_path}'.") from exc
+
+
+def get_method_class(method_name: str, dataset_name: str = None):
+    """Return the *class* corresponding to *method_name* (case-insensitive).  This
+    function also performs a fuzzy match when *dataset_name* is provided, to
+    pick the variant tailored for that dataset (e.g. agentverse_humaneval).
+    """
     method_name = method_name.lower()
 
-    all_method_names = method2class.keys()
-    matched_method_names = [sample_method_name for sample_method_name in all_method_names if method_name in sample_method_name]
-
-    if len(matched_method_names) > 0:
-        if dataset_name is not None:
-            # lowercase the dataset name
-            dataset_name = dataset_name.lower()
-            # check if there are method names that contain the dataset name
-            matched_method_data_names = [sample_method_name for sample_method_name in matched_method_names if sample_method_name.split('_')[-1] in dataset_name]
-            if len(matched_method_data_names) > 0:
-                method_name = matched_method_data_names[0]
-                if len(matched_method_data_names) > 1:
-                    print(f"[WARNING] Found multiple methods matching {dataset_name}: {matched_method_data_names}. Using {method_name} instead.")
-    else:
-        raise ValueError(f"[ERROR] No method found matching {method_name}. Please check the method name.")
+    # First stage: direct match or substring match within registry keys
+    candidate_keys = [key for key in _method_registry if method_name in key]
+    if not candidate_keys:
+        raise ValueError(f"[ERROR] No method found matching '{method_name}'. Available: {_method_registry.keys()}")
+
+    # If dataset name is provided, further filter by suffix match
+    if dataset_name:
+        dataset_name = dataset_name.lower()
+        ds_filtered = [key for key in candidate_keys if key.split("_")[-1] in dataset_name]
+        if ds_filtered:
+            candidate_keys = ds_filtered
+
+    # Use first candidate if multiple remain (warn user)
+    chosen_key = candidate_keys[0]
+    if len(candidate_keys) > 1:
+        print(f"[WARNING] Multiple methods matched '{method_name}': {candidate_keys}. Using '{chosen_key}'.")
+
+    module_path, class_name = _method_registry[chosen_key]
+
+    try:
+        method_cls = _import_class(module_path, class_name)
+    except Exception as exc:
+        raise ImportError(
+            f"Failed to import method '{chosen_key}' (module '{module_path}', class '{class_name}'). "
+            "Please ensure all optional dependencies are installed." ) from exc
 
-    return method2class[method_name]
+    return method_cls

methods/maas/.maas/config2.yaml

@@ -0,0 +1,14 @@
+# Full Example: https://github.com/geekan/MetaGPT/blob/main/config/config2.example.yaml
+# Reflected Code: https://github.com/geekan/MetaGPT/blob/main/metagpt/config2.py
+# Config Docs: https://docs.deepwisdom.ai/main/en/guide/get_started/configuration.html
+llm:
+  api_type: "openai"  # or azure / ollama / groq etc.
+  model: "gpt-4o-mini"  #  gpt-3.5-turbo or gpt-4-turbo
+  base_url: ""  # or forward url / other llm url
+  api_key: ""
+models:
+  gpt-4o-mini:
+    api_type: "openai"  # or azure / ollama / groq etc.
+    model: "gpt-4o-mini"  #  gpt-3.5-turbo or gpt-4-turbo
+    base_url: ""  # or forward url / other llm url
+    api_key: ""

methods/maas/MASLAB_INTEGRATION.md

@@ -0,0 +1,153 @@
+# MAAS Integration in MASLab
+
+## Overview
+
+MAAS (Multi-agent Architecture Search via Agentic Supernet) has been integrated into MASLab. This integration allows you to use MAAS's advanced multi-agent architecture search capabilities within the MASLab framework.
+
+## Key Parameters
+
+### Round and Sample Parameters
+- **round**: Specifies which training round's model to use (default: 1)
+- **sample**: Specifies which sample size model to use (default: 4)
+
+Pre-trained models are stored in:
+```
+methods/maas/maas/ext/maas/scripts/optimized/{DATASET}/train/round_{round}/{DATASET}_controller_sample{sample}.pth
+```
+
+For example, with round=1 and sample=4 for MATH dataset:
+```
+methods/maas/maas/ext/maas/scripts/optimized/MATH/train/round_1/MATH_controller_sample4.pth
+```
+
+## Usage
+
+### 1. Inference Mode (Default)
+
+To run inference with a pre-trained MAAS model:
+
+```bash
+# Using default round=1, sample=4
+python inference.py --method_name maas --test_dataset_name MATH --debug
+
+# Or specify custom round and sample in config
+python inference.py --method_name maas --test_dataset_name MATH --method_config_name config_custom --debug
+```
+
+**Note**: Ensure the pre-trained model file exists for your specified round and sample parameters.
+
+### 2. Training Mode
+
+To train a MAAS model:
+
+```bash
+# First, build your dataset
+python datasets/build_test_dataset.py --dataset_name MATH
+
+# Then run training
+python inference.py --method_name maas --test_dataset_name MATH --method_config_name config_train --require_val
+```
+
+### 3. Batch Inference on Dataset
+
+For full dataset evaluation:
+
+```bash
+# Sequential processing
+python inference.py --method_name maas --test_dataset_name MATH --sequential
+
+# Parallel processing
+python inference.py --method_name maas --test_dataset_name MATH
+```
+
+## Configuration Options
+
+Create custom configuration files in `methods/maas/configs/`:
+
+```yaml
+# Example: config_custom.yaml
+is_test: true      # true for inference, false for training
+sample: 4          # Must match your pre-trained model
+round: 1           # Must match your pre-trained model
+batch_size: 4      # Batch size for training
+lr: 0.01           # Learning rate
+is_textgrad: false # Whether to use textgrad optimization
+```
+
+### Pre-configured Files:
+- `config_main.yaml`: Default inference configuration (round=1, sample=4)
+- `config_train.yaml`: Training configuration
+
+## Supported Datasets
+
+MAAS currently supports:
+- **MATH**: Mathematical problem solving
+- **GSM8K**: Grade school math problems
+- **HumanEval**: Code generation tasks
+
+## Data Preparation
+
+The integration automatically handles data conversion:
+1. MASLab JSON format → MAAS JSONL format
+2. Automatic copying to MAAS data directory
+3. Both train and test splits handled
+
+Data flow:
+```
+datasets/data/{dataset}.json → methods/maas/maas/ext/maas/data/{dataset}_{train|test}.jsonl
+```
+
+## Pre-trained Models
+
+### Available Models
+Currently available pre-trained model:
+- MATH dataset: round_1, sample_4
+
+### Using Pre-trained Models
+1. Ensure your config matches the available model (round and sample)
+2. The model file must exist in the expected path
+3. If model not found, you'll see an error message with the expected path
+
+## Example Workflows
+
+### Quick Test with Pre-trained Model
+```bash
+# Test on MATH dataset with existing model (round=1, sample=4)
+python inference.py --method_name maas --test_dataset_name MATH --model_name gpt-4o-mini --debug
+```
+
+### Train New Model
+```bash
+# 1. Prepare dataset
+python datasets/build_test_dataset.py --dataset_name GSM8K
+
+# 2. Train with custom parameters
+# Create config_gsm8k_train.yaml with desired parameters
+python inference.py --method_name maas --test_dataset_name GSM8K --method_config_name config_gsm8k_train --require_val
+```
+
+## Troubleshooting
+
+### "No pre-trained model found"
+- Check if the model file exists in the expected path
+- Verify round and sample parameters in your config
+- Train a model first if needed
+
+### API Configuration Issues
+- Ensure your model API configuration is properly set
+- Check API keys and endpoints
+- Use a valid model configuration file
+
+## Limitations
+
+1. **Single Sample Inference**: MAAS is designed for batch processing. Single sample inference returns a simplified response.
+2. **Model Dependency**: Inference requires pre-trained models with matching round and sample parameters.
+3. **Training Time**: MAAS training can be computationally intensive due to architecture search.
+
+## Future Improvements
+
+1. Implement proper single-sample inference support
+2. Add automatic model downloading/sharing
+3. Support for more datasets
+4. Integrate MAAS's visualization capabilities
+5. Support for distributed training

methods/maas/__init__.py

@@ -0,0 +1,24 @@
+"""
+MASLab integration for the maas multi-agent framework.
+
+This module sets up the import paths so that the maas package can be imported
+correctly within the MASLab structure without modifying all internal imports.
+"""
+
+import sys
+import os
+from pathlib import Path
+
+# Add the maas package to Python path so internal absolute imports work
+maas_package_dir = Path(__file__).parent
+if str(maas_package_dir) not in sys.path:
+    sys.path.insert(0, str(maas_package_dir))
+
+# Apply config patch BEFORE importing any maas modules
+from .maas_config_patch import apply_config_patch
+apply_config_patch()
+
+# Import the main class
+from .maas_main import MAAS_Main
+
+__all__ = ['MAAS_Main']