SLMaGNN

Reference implementation for the JUCS article SLMaGNN: A Small Language Model-Assisted Framework for Exploratory Graph Neural Network Design.

The code reproduces the iterative SLM-in-the-loop search described in the paper: a local Ollama model proposes GNN configurations, a validator filters invalid architectures, each valid candidate is trained and evaluated, and the resulting metrics are fed back into the next iteration.

Directory layout

code/
├── slmagnn.py                  # CLI entry point
├── slmagnn/                    # Python package
│   ├── constants.py            # Search-space constants
│   ├── data.py                 # Dataset loading and splits
│   ├── search_space.py         # Parsing, validation, sampling
│   ├── model.py                # GNN construction
│   ├── training.py             # Training loop helpers
│   ├── evaluation.py           # Metrics and checkpoint evaluation
│   ├── llm.py                  # Ollama / structured SLM proposals
│   └── runner.py               # Main search loop
├── baselines/
│   ├── tpe_search.py           # TPE baseline (Optuna)
│   └── grid_search.py          # Grid-search baseline
├── scripts/
│   ├── setup_environment.sh    # Virtualenv + pip install
│   └── pull_ollama_models.sh   # Download Ollama models
├── experiment_config.env.example
├── .env.example
├── requirements.txt
├── requirements-baselines.txt
└── runner_methodology.sh       # Batch campaign launcher

Requirements

• Python 3.10+
• PyTorch and PyTorch Geometric (install for your CUDA/CPU setup)
• Ollama for local SLM inference
• Optional: Optuna for the TPE baseline

Quick start

cd code
./scripts/setup_environment.sh
source .venv/bin/activate

# Install PyTorch / PyG for your platform, then:
ollama serve
./scripts/pull_ollama_models.sh

cp .env.example .env
cp experiment_config.env.example experiment_config.env

Smoke test (SLM search)

python slmagnn.py \
  --dataset Cora \
  --model codegemma \
  --search_mode slm_search \
  --iterations 5 \
  --epochs 20 \
  --selection_split val \
  --report_split test \
  --skip_plots

Random-search baseline

python slmagnn.py \
  --dataset Cora \
  --search_mode random_search \
  --iterations 50 \
  --epochs 100 \
  --skip_plots

Full paper campaign

Edit experiment_config.env, then:

./runner_methodology.sh --config experiment_config.env

Environment variables

Variable	Default	Description
OLLAMA_BASE_URL	http://localhost:11434/v1	OpenAI-compatible Ollama endpoint
OLLAMA_API_KEY	ollama	Placeholder API key required by the OpenAI client

These can be set in .env or exported in the shell.

Search modes

Mode	CLI value	Description
SLM search	slm_search	Ollama-backed architecture proposals
Random search	random_search	Uniform sampling from the shared search space
TPE	tpe_search	Optuna TPE baseline
Grid search	grid_search	Finite grid baseline

Paper defaults: ITERATIONS=500, EPOCHS=300, STAGNATION_LIMIT=30, EARLY_STOPPING_PATIENCE=20, validation-based selection and held-out test reporting.

Evaluated Ollama models

codegemma, codellama, codeqwen, codestral, gemma:2b, gemma:7b, llama3, mistral, opencoder:1.5b, opencoder:8b, phi3

Datasets

Cora, CiteSeer, PubMed, WikiCS, Chameleon, Squirrel, Actor, Cornell, Texas, Wisconsin

Datasets are downloaded automatically by PyTorch Geometric into data/.

Outputs

Each run writes to:

results/runs/<dataset>/<method>/<model>/seed_<seed>/
├── campaign.json
├── history.json
├── metrics.json
├── best_model.pt
└── run.log

Aggregated rows are appended to results/tables/all_runs.csv.

Citation

If you use this code, please cite the SLMaGNN JUCS article (authors: Krzywda, Łukasik, Gandomi).