echo / code /scripts /run_generate_metadata.sh
amonshano's picture
Add Echo-Memory codebase used for this run (CC BY 4.0, JD Echo Team) (part 4)
c335050 verified
Raw
History Blame Contribute Delete
2.33 kB
#!/bin/bash
# Generate metadata CSV for a context-based memory dataset.
set -euo pipefail
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
REPO_ROOT="${REPO_ROOT:-$(cd "${SCRIPT_DIR}/.." && pwd)}"
cd "${REPO_ROOT}"
DETECTED_CPUS=$(python3 -c "import os; print(os.cpu_count())")
OPTIMAL_WORKERS=$((DETECTED_CPUS > 4 ? DETECTED_CPUS - 2 : DETECTED_CPUS))
NUM_WORKERS=${NUM_WORKERS:-$OPTIMAL_WORKERS}
DATASET_BASE_PATH="${DATASET_BASE_PATH:-${REPO_ROOT}/data/Context-as-Memory-Dataset}"
OUTPUT_CSV="${OUTPUT_CSV:-${DATASET_BASE_PATH}/metadata_full.csv}"
SEGMENT_LENGTH="${SEGMENT_LENGTH:-81}"
CONTEXT_FRAMES="${CONTEXT_FRAMES:-5}"
METADATA_MAX_ROWS="${METADATA_MAX_ROWS:-${DATASET_SIZE_ROWS:-0}}"
echo "=========================================="
echo "Context-as-Memory Metadata Generation"
echo "=========================================="
echo "Detected CPUs: $DETECTED_CPUS"
echo "Using workers: $NUM_WORKERS"
echo "Dataset: ${DATASET_BASE_PATH}"
echo "Output: ${OUTPUT_CSV}"
echo "Max rows: ${METADATA_MAX_ROWS} (0 = full metadata)"
echo "=========================================="
echo ""
python3 src/data/preprocess_cam_dataset.py \
--dataset_base_path "${DATASET_BASE_PATH}" \
--output_csv "${OUTPUT_CSV}" \
--segment_length "${SEGMENT_LENGTH}" \
--context_frames "${CONTEXT_FRAMES}"
if [ "${METADATA_MAX_ROWS}" != "0" ]; then
python3 - "${OUTPUT_CSV}" "${METADATA_MAX_ROWS}" <<'PY'
import csv
import os
import sys
path = sys.argv[1]
max_rows = int(sys.argv[2])
if max_rows < 0:
raise SystemExit("METADATA_MAX_ROWS must be >= 0")
if max_rows > 0:
tmp_path = f"{path}.tmp"
with open(path, newline="", encoding="utf-8") as src, open(tmp_path, "w", newline="", encoding="utf-8") as dst:
reader = csv.reader(src)
writer = csv.writer(dst)
header = next(reader, None)
if header is not None:
writer.writerow(header)
for idx, row in enumerate(reader):
if idx >= max_rows:
break
writer.writerow(row)
os.replace(tmp_path, path)
print(f"Truncated metadata to {max_rows} rows: {path}")
PY
fi
echo ""
echo "Generation complete! Checking results..."
if [ -f "${OUTPUT_CSV}" ]; then
echo "CSV file generated successfully"
wc -l "${OUTPUT_CSV}"
else
echo "CSV file generation failed"
fi