scripts/0_download_data.py

import os
import sys
import shutil
from pathlib import Path
import dotenv

sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

from utils.logger import setup_logger
from utils.json import get_file_path_from_config

log = setup_logger(__name__)

fast_text_bin_path = get_file_path_from_config("embeddings.fasttext_bin")
fast_text_vec_path = get_file_path_from_config("embeddings.fasttext_vec")


def download_embeddings() -> None:
    """Hugging Face Datasets から Embeddings をダウンロードする"""
    target_files = {
        "cc.ja.300.bin": Path(fast_text_bin_path),
        "cc.ja.300.vec": Path(fast_text_vec_path),
    }

    missing_files = [name for name, path in target_files.items() if not path.exists()]
    if not missing_files:
        log.info("Embeddings already exist. Skipping download.")
        return

    from huggingface_hub import hf_hub_download

    # フォルダがなかったら新規作成
    target_dir = Path(fast_text_bin_path).parent
    os.makedirs(target_dir, exist_ok=True)
    # huggingface_hub は local_dir 配下に .cache/huggingface を作るため先に用意
    try:
        os.makedirs(target_dir / ".cache" / "huggingface", exist_ok=True)
    except PermissionError:
        log.error(
            "権限エラー: %s に .cache/huggingface を作成できません。Dockerfile の権限設定を確認してください。",
            target_dir,
        )
        raise

    log.info("Downloading embeddings from Hugging Face Datasets...")

    dotenv.load_dotenv()  # .envから環境変数をロード

    # ダウンロード
    if "cc.ja.300.bin" not in missing_files:
        log.info("Embeddings binary already exists. Skipping download.")
    else:
        path = hf_hub_download(
            repo_id=os.getenv("HF_EMBEDDINGS_REPO_ID"),
            repo_type="dataset",
            filename="cc.ja.300.bin",
            local_dir=str(target_dir),
            token=os.getenv("HF_TOKEN"),
        )
        log.info(f"Embeddings downloaded: {path}")
        _ensure_in_target_dir("cc.ja.300.bin", str(target_dir))

    if "cc.ja.300.vec" not in missing_files:
        log.info("Embeddings vector already exists. Skipping download.")
    else:
        path = hf_hub_download(
            repo_id=os.getenv("HF_EMBEDDINGS_REPO_ID"),
            repo_type="dataset",
            filename="cc.ja.300.vec",
            local_dir=str(target_dir),
            token=os.getenv("HF_TOKEN"),
        )
        log.info(f"Embeddings downloaded: {path}")
        _ensure_in_target_dir("cc.ja.300.vec", str(target_dir))

    return


def _ensure_in_target_dir(filename: str, target_dir: str) -> None:
    """
    ダウンロードされたファイルが target_dir 直下に無い場合、
    target_dir 配下（例: `.cache/huggingface`）から探して移動する。
    """
    dest = Path(target_dir) / filename
    if dest.exists():
        return

    # 既知のサブディレクトリ優先
    candidates: list[Path] = []
    sub = Path(target_dir) / ".cache" / "huggingface"
    if sub.exists():
        candidates.extend(p for p in sub.rglob(filename))

    # fallback: target_dir 以下を再帰探索
    if not candidates:
        candidates = list(Path(target_dir).rglob(filename))

    if not candidates:
        log.warning("期待するファイルが見つかりません: %s", filename)
        return

    src = candidates[0]
    try:
        dest.parent.mkdir(parents=True, exist_ok=True)
        shutil.move(str(src), str(dest))
        log.info("Moved %s -> %s", src, dest)
    except Exception as e:
        log.error("%s の移動に失敗しました: %s", filename, e)
        raise


if __name__ == "__main__":
    download_embeddings()