from huggingface_hub import hf_hub_download
import os

repo_id = "bowen-upenn/PersonaMem"
local_dir = "data/raw_datasets/personamem"
files_to_download = [
    "questions_32k.csv",
    "shared_contexts_32k.jsonl"
]

os.makedirs(local_dir, exist_ok=True)

print(f"Downloading files from {repo_id} to {local_dir}...")

for filename in files_to_download:
    print(f"Downloading {filename}...")
    hf_hub_download(
        repo_id=repo_id,
        filename=filename,
        repo_type="dataset",
        local_dir=local_dir
    )

print("Download complete.")