repro_bundle/scripts/run_pathstar.sh

#!/bin/bash
set -euxo pipefail

pip install -q tqdm networkx umap-learn scikit-learn matplotlib

cd /workspace/repo

STAR_DEGREE=${STAR_DEGREE:-10000}
PATH_LENGTH=${PATH_LENGTH:-6}
EDGE_EPOCHS=${EDGE_EPOCHS:-2500}
PATH_EPOCHS=${PATH_EPOCHS:-10000}
TAG=${TAG:-run}

python data/build_in_weights_datasets.py \
  --graph_type star --star_degree ${STAR_DEGREE} --star_subtree_degree 1 \
  --path_length ${PATH_LENGTH} --add_forward_edges --add_backward_edges --random_seed 0

python train_in_weights.py \
  --training_recipe mixed_full_path --model_family gpt \
  --graph_type star --star_degree ${STAR_DEGREE} --star_subtree_degree 1 \
  --path_length ${PATH_LENGTH} \
  --add_forward_edges --add_backward_edges \
  --edge_memorization_epochs ${EDGE_EPOCHS} --path_finetuning_epochs ${PATH_EPOCHS} \
  --edge_memorization_eval_interval_epochs 100 --path_finetuning_eval_interval_epochs 100 \
  --track_embedding_evolution \
  --experiment_log_root /data/experiment_logs_${TAG} 2>&1 | tee /data/train_log_${TAG}.txt

python /workspace/analysis/analyze_run.py --experiment_log_root /data/experiment_logs_${TAG} --tag ${TAG} --out_dir /data/analysis_${TAG}

echo "JOB_DONE_${TAG}"