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()