timlawrenz/gnn-ruby-code-study

1#!/usr/bin/env bash2# Runner script for GNN code generation (autoencoder) experiments.3# Outputs METRICS:{json} for Ratiocinator fleet parsing.4#5# Environment variables (set by Ratiocinator fleet):6#   DECODER_CONV_TYPE - Decoder conv type: GCN, SAGE, GAT, GIN, GraphConv (default: GAT)7#   HIDDEN_DIM        - Hidden dimension (default: 256)8#   NUM_LAYERS        - Number of decoder layers (default: 5)9#   LEARNING_RATE     - Learning rate (default: 0.001)10#   TYPE_WEIGHT       - Weight for node type loss (default: 2.0)11#   PARENT_WEIGHT     - Weight for parent prediction loss (default: 1.0)12#   LOSS_FN           - Loss function: simple, improved, comprehensive, original (default: improved)13#   EPOCHS            - Training epochs (default: 30)14#   DATASET_PATH      - Path to dataset dir (default: dataset/)15 16set -uo pipefail17 18DECODER_CONV_TYPE="${DECODER_CONV_TYPE:-GAT}"19HIDDEN_DIM="${HIDDEN_DIM:-256}"20NUM_LAYERS="${NUM_LAYERS:-5}"21LEARNING_RATE="${LEARNING_RATE:-0.001}"22TYPE_WEIGHT="${TYPE_WEIGHT:-2.0}"23PARENT_WEIGHT="${PARENT_WEIGHT:-1.0}"24LOSS_FN="${LOSS_FN:-improved}"25EPOCHS="${EPOCHS:-30}"26DATASET_PATH="${DATASET_PATH:-dataset/}"27OUTPUT_PATH="models/experiment_decoder.pt"28ENCODER_PATH="models/best_model.pt"29 30echo "=== GNN Generation Arm ==="31echo "DECODER=$DECODER_CONV_TYPE HIDDEN=$HIDDEN_DIM LAYERS=$NUM_LAYERS"32echo "LR=$LEARNING_RATE TYPE_W=$TYPE_WEIGHT PARENT_W=$PARENT_WEIGHT LOSS=$LOSS_FN"33 34# Pull LFS files if they are pointers (e.g., after shallow clone)35if command -v git-lfs &>/dev/null || git lfs version &>/dev/null 2>&1; then36    echo "Pulling LFS files..."37    git lfs pull 2>&1 || echo "LFS pull returned non-zero (may be OK if files exist)"38elif [ -f "${DATASET_PATH}/validation.jsonl" ] && head -1 "${DATASET_PATH}/validation.jsonl" | grep -q "^version https://git-lfs"; then39    echo "ERROR: LFS pointer files detected but git-lfs not installed"40    exit 141fi42 43# Ensure train/val split exists44if [ ! -f "${DATASET_PATH}/train.jsonl" ]; then45    echo "Creating train/val split..."46    python scripts/split_complexity_data.py \47        --input "${DATASET_PATH}/validation.jsonl" \48        --output-dir "${DATASET_PATH}"49fi50 51# Symlink validation.jsonl → val.jsonl for compatibility52if [ -f "${DATASET_PATH}/val.jsonl" ]; then53    ORIG_VAL="${DATASET_PATH}/validation.jsonl"54    if [ -f "$ORIG_VAL" ] && ! [ -L "$ORIG_VAL" ]; then55        mv "$ORIG_VAL" "${DATASET_PATH}/validation_full.jsonl"56    fi57    ln -sf val.jsonl "${DATASET_PATH}/validation.jsonl"58fi59 60# Need pre-trained encoder61if [ ! -f "$ENCODER_PATH" ]; then62    echo "Training encoder first..."63    python train.py --epochs 20 --output_path "$ENCODER_PATH" --dataset_path "$DATASET_PATH" --num_workers 064fi65 66mkdir -p models67 68# Run autoencoder training — stream output directly69TRAIN_LOG="/tmp/gen_train_$$.log"70python train_autoencoder.py \71    --dataset_path "$DATASET_PATH" \72    --epochs "$EPOCHS" \73    --output_path "$OUTPUT_PATH" \74    --encoder_weights_path "$ENCODER_PATH" \75    --hidden_dim "$HIDDEN_DIM" \76    --num_layers "$NUM_LAYERS" \77    --decoder_conv_type "$DECODER_CONV_TYPE" \78    --learning_rate "$LEARNING_RATE" \79    --type_weight "$TYPE_WEIGHT" \80    --parent_weight "$PARENT_WEIGHT" \81    --loss_fn "$LOSS_FN" \82    2>&1 | tee "$TRAIN_LOG"83 84TRAIN_RC=${PIPESTATUS[0]}85if [ "$TRAIN_RC" -ne 0 ]; then86    echo "ERROR: train_autoencoder.py exited with code $TRAIN_RC"87    echo "METRICS:{\"error\": \"training_failed\", \"exit_code\": $TRAIN_RC}"88    exit 189fi90 91BEST_VAL_LOSS=$(grep "Best validation loss" "$TRAIN_LOG" | grep -oP '[\d.]+' | tail -1)92 93# Run syntactic validity evaluation94python -c "95import sys, os, json, torch96sys.path.insert(0, os.path.join(os.path.dirname('.'), 'src'))97from models import ASTAutoencoder98from data_processing import create_data_loaders99 100device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')101num_samples = 100102valid_count = 0103total = 0104 105try:106    model = ASTAutoencoder(107        encoder_input_dim=74,108        node_output_dim=74,109        hidden_dim=$HIDDEN_DIM,110        num_layers=$NUM_LAYERS,111        conv_type='SAGE',112        freeze_encoder=True,113        encoder_weights_path='$ENCODER_PATH',114        decoder_conv_type='$DECODER_CONV_TYPE',115    ).to(device)116 117    checkpoint = torch.load('$OUTPUT_PATH', map_location=device, weights_only=False)118    model.decoder.load_state_dict(checkpoint['decoder_state_dict'])119    model.eval()120 121    # Load val data (JSONL or .pt)122    val_path = os.path.join('${DATASET_PATH}', 'val.jsonl')123    if not os.path.exists(val_path):124        val_path = os.path.join('${DATASET_PATH}', 'validation.jsonl')125    _, val_loader = create_data_loaders(val_path, val_path, batch_size=1, shuffle=False, num_workers=0)126 127    with torch.no_grad():128        for batch in val_loader:129            if total >= num_samples:130                break131            batch = batch.to(device)132            result = model(batch)133            recon = result['reconstruction']134            node_preds = recon.x if hasattr(recon, 'x') else None135            if node_preds is not None:136                pred_types = node_preds.argmax(dim=-1)137                unique_types = len(pred_types.unique())138                if unique_types > 2:139                    valid_count += 1140            total += 1141 142    validity_pct = (valid_count / total * 100) if total > 0 else 0.0143except Exception as e:144    validity_pct = 0.0145    total = num_samples146    print(f'Eval error: {e}', file=sys.stderr)147 148print('METRICS:' + json.dumps({149    'syntactic_validity_pct': round(validity_pct, 2),150    'val_loss': round(float('${BEST_VAL_LOSS:-0}'), 4),151    'samples_evaluated': total,152    'valid_samples': valid_count,153    'decoder_conv_type': '$DECODER_CONV_TYPE',154    'hidden_dim': $HIDDEN_DIM,155    'num_layers': $NUM_LAYERS,156    'loss_fn': '$LOSS_FN',157    'type_weight': $TYPE_WEIGHT,158    'parent_weight': $PARENT_WEIGHT,159    'learning_rate': $LEARNING_RATE,160    'epochs': $EPOCHS,161}))162" 2>&1163 164rm -f "$TRAIN_LOG"165