social-data-attribution-w2/scripts/modal/run_materialize_sample.sh

#!/usr/bin/env bash
set -euo pipefail

# Materialize a sample and upload to R2, HF Buckets, and/or HF Datasets.
#
# Usage:
#   bash scripts/modal/run_materialize_sample.sh data/samples/sample_500_docs
#   bash scripts/modal/run_materialize_sample.sh data/samples/sample_500_docs 128 --upload-only
#   UPLOAD_TO=hf-buckets bash scripts/modal/run_materialize_sample.sh data/samples/sample_500_docs
#   UPLOAD_TO=r2,hf-buckets bash scripts/modal/run_materialize_sample.sh data/samples/sample_500_docs
#   UPLOAD_TO=r2,hf bash scripts/modal/run_materialize_sample.sh data/samples/sample_500_docs

SAMPLE_DIR="${1:?Usage: run_materialize_sample.sh <sample_dir> [chunk_count] [--upload-only]}"
CHUNK_COUNT="${2:-128}"
HF_ORG="${HF_ORG:-HCAI-Lab}"
UPLOAD_TO="${UPLOAD_TO:-hf-buckets}"

EXTRA_FLAGS=""
for arg in "$@"; do
    if [ "$arg" = "--upload-only" ]; then
        EXTRA_FLAGS="$EXTRA_FLAGS --upload-only"
    fi
done

SAMPLE_NAME=$(basename "$SAMPLE_DIR")
VOLUME_NAME="soc134-output-$(echo "$SAMPLE_NAME" | tr '_' '-')"
HF_REPO="${HF_ORG}/dolma3-6t-${SAMPLE_NAME}"
HF_BUCKET_ID="${HF_BUCKET_ID:-${HF_ORG}/dolma3-6t-$(echo "$SAMPLE_NAME" | tr '_' '-')}"
MANIFEST="${SAMPLE_DIR}/working_sample_manifest.parquet"

if [ ! -f "$MANIFEST" ]; then
    echo "ERROR: manifest not found: $MANIFEST"
    exit 1
fi

echo "Sample:    $SAMPLE_NAME"
echo "Volume:    $VOLUME_NAME"
echo "Upload to: $UPLOAD_TO"
echo "Chunks:    $CHUNK_COUNT"
if echo "$UPLOAD_TO" | grep -q "hf-buckets"; then
    echo "HF bucket: $HF_BUCKET_ID"
elif echo "$UPLOAD_TO" | grep -q "hf"; then
    echo "HF repo:   $HF_REPO"
fi
if echo "$UPLOAD_TO" | grep -q "r2"; then
    echo "R2 prefix: soc134-samples/$SAMPLE_NAME/"
fi
echo ""

SOC134_OUTPUT_VOLUME="$VOLUME_NAME" \
    uv run modal run scripts/modal/materialize_working_sample.py \
        --sample-manifest "$MANIFEST" \
        --sample-name "$SAMPLE_NAME" \
        --hf-repo "$HF_REPO" \
        --hf-bucket-id "$HF_BUCKET_ID" \
        --chunk-count "$CHUNK_COUNT" \
        --upload-to "$UPLOAD_TO" \
        $EXTRA_FLAGS