Blame for plus/kaggle_plus.py

"""Kaggle: демонстрация и прогон пакета PLUS на своей пачке треков.

Что делает ячейка:
  1. ставит зависимости (pretty_midi, librosa, soundfile, scipy);
  2. скачивает пакет с Hugging Face (или берёт локальную копию);
  3. прогоняет самопроверку на синтетике: чистая правда -> три «движка»
     с ошибками -> слияние -> F1, чтобы видеть, что пакет реально работает;
  4. обрабатывает папку AUDIO_DIR (если в ней есть MIDI) и сводит отчёт;
  5. по флагу заливает результаты на Hugging Face.

Локальный запуск:  python3 kaggle_plus.py --demo
Своя папка:       python3 kaggle_plus.py --input-dir /kaggle/input/takes --out-dir /kaggle/working/plus
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys

PACK = "plus.zip"
REPO = "Nokia123543/muscriptor-large-fork"
WORK = "/kaggle/working" if os.path.isdir("/kaggle/working") else os.getcwd()


def step(number: int, total: int, text: str) -> None:
    print("[%d/%d] %s" % (number, total, text), flush=True)


def ensure_deps() -> None:
    need = []
    for module, package in (("pretty_midi", "pretty_midi"), ("librosa", "librosa"),
                            ("soundfile", "soundfile"), ("scipy", "scipy")):
        try:
            __import__(module)
        except Exception:
            need.append(package)
    if need:
        step(1, 5, "ставлю: " + ", ".join(need))
        subprocess.run([sys.executable, "-m", "pip", "install", "-q"] + need, check=True)


def ensure_pack(target: str) -> str:
    here = os.path.dirname(os.path.abspath(__file__))
    if os.path.isfile(os.path.join(here, "plus_core.py")):
        step(2, 5, "пакет найден рядом: " + here)
        return here
    bundle = os.path.join(target, "plus")
    if not os.path.isdir(bundle):
        step(2, 5, "скачиваю пакет с HF")
        os.makedirs(target, exist_ok=True)
        from huggingface_hub import hf_hub_download
        archive = hf_hub_download(repo_id=REPO, filename=PACK, local_dir=target)
        import zipfile
        with zipfile.ZipFile(archive) as zf:
            zf.extractall(target)
    step(2, 5, "пакет распакован: " + bundle)
    return bundle


def demo(pack_dir: str) -> dict:
    """Синтетика: правда, три движка с ошибками, слияние, метрики."""
    sys.path.insert(0, pack_dir)
    import numpy as np
    from plus_core import json_default, Note, Piece, save_midi
    from plus_ensemble import compare, merge
    from plus_transcribe import run_pipeline

    step(3, 5, "самопроверка: собираю правду и три «движка»")
    rng = np.random.default_rng(11)
    beat = 60 / 72 / 2
    notes = []
    melody = [72, 74, 76, 77, 79, 77, 76, 74] * 4
    for index, pitch in enumerate(melody):
        notes.append(Note(max(0.0, index * beat + rng.normal(0, 0.010)), beat * 0.9, pitch, 90))
    for bar in range(32):
        root, shape = [(48, (0, 4, 7)), (53, (0, 3, 7)), (55, (0, 4, 7)), (48, (0, 4, 7))][bar % 4]
        start = max(0.0, bar * beat * 2 + rng.normal(0, 0.006))
        for offset in shape:
            notes.append(Note(start, beat * 1.8, root + offset, 80))
    truth = Piece(sorted(notes, key=lambda n: n.start), [(0.0, 72.0)])

    work = os.path.join(WORK, "plus_demo")
    os.makedirs(work, exist_ok=True)
    truth_path = os.path.join(work, "truth.mid")
    save_midi(truth, truth_path)

    sources = []
    for tag, seed, weight in (("muscriptor", 1, 1.0), ("aria", 2, 1.0)):
        r = np.random.default_rng(seed)
        variant = [Note(**vars(n)) for n in truth.notes if r.random() > 0.10]
        for _ in range(4):
            variant.append(Note(float(r.uniform(0, 20)), 0.30, int(r.integers(50, 88)), 90))
        variant.sort(key=lambda n: n.start)
        path = os.path.join(work, "engine_%s.mid" % tag)
        save_midi(Piece(variant, [(0.0, 72.0)]), path)
        sources.append((tag, variant, weight))

    step(4, 5, "слияние голосованием")
    merged, stats = merge(sources, tolerance=0.06, min_votes=2.0, keep_strong=False, tempo=72.0)
    merged_path = os.path.join(work, "merged.mid")
    save_midi(merged, merged_path)

    result = {"merged_f1": compare(truth.notes, merged.notes),
              "sources": {name: compare(truth.notes, notes_) for name, notes_, _w in sources},
              "agreement": stats}

    step(5, 5, "пост-обработка и рендер превью")
    polished = merged.copy()
    pipeline_report = run_pipeline(polished, ["cleanup", "quantize", "hands", "pedal",
                                             "chords", "report"], {})
    save_midi(polished, os.path.join(work, "merged_plus.mid"))
    subprocess.run([sys.executable, os.path.join(pack_dir, "plus_render.py"),
                    "--midi", os.path.join(work, "merged_plus.mid"),
                    "--out", os.path.join(work, "preview.wav"), "--stop", "12"], check=False)

    best = max(result["sources"].items(), key=lambda kv: kv[1]["f1"])
    print("\nF1 слияния %.3f | лучший вход %.3f (%s) | согласие %.2f источника"
          % (result["merged_f1"]["f1"], best[1]["f1"], best[0],
             result["agreement"]["agreement_mean"]), flush=True)
    print("файлы: " + work, flush=True)
    result["after_pipeline_f1"] = compare(truth.notes, polished.notes)
    result["pipeline"] = polished.stats()
    result["pipeline_stages"] = {k: v for k, v in pipeline_report["stages"].items()
                                 if k in ("cleanup", "quantize", "hands", "pedal", "chords")}
    with open(os.path.join(work, "report.json"), "w", encoding="utf-8") as fh:
        json.dump(result, fh, ensure_ascii=False, indent=1, default=str)
    return result



def _hf_token() -> str | None:
    """Токен из окружения или из Kaggle Secrets, если ячейка запущена в Kaggle."""
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_TOKEN")
    if token:
        return token
    try:
        from kaggle_secrets import UserSecretsClient
        return UserSecretsClient().get_secret("HF_TOKEN")
    except Exception:
        return None


def _files(folder: str, with_audio: bool) -> list[str]:
    if not os.path.isdir(folder):
        return []
    skip = () if with_audio else (".wav", ".flac", ".mp3")
    return [os.path.join(folder, n) for n in sorted(os.listdir(folder))
            if os.path.isfile(os.path.join(folder, n)) and not n.lower().endswith(skip)]


def upload_build(folders: dict[str, str], repo: str, stamp: str,
                 with_audio: bool = True, token: str | None = None) -> dict:
    """Заливает сборку на HF в builds/<метка времени>/<имя папки>.

    Возвращает словарь с числом файлов, объёмом и ссылками. Ничего не делает,
    если токена нет: сборка остаётся локально, и это печатается.
    """
    token = token or _hf_token()
    if not token:
        print("заливка пропущена: нет HF_TOKEN (в Kaggle это Add-ons > Secrets)", flush=True)
        return {"uploaded": 0, "reason": "no_token"}

    from huggingface_hub import HfApi
    api = HfApi(token=token)
    uploaded, total_bytes = 0, 0
    links: list[str] = []
    for name, folder in folders.items():
        chosen = _files(folder, with_audio)
        if not chosen:
            continue
        target = "builds/%s/%s" % (stamp, name)
        print("[заливка] %s -> %s (%d файлов)" % (folder, target, len(chosen)), flush=True)
        api.upload_folder(folder_path=folder, path_in_repo=target, repo_id=repo,
                          allow_patterns=[os.path.basename(f) for f in chosen],
                          commit_message="build %s: %s" % (stamp, name))
        for path in chosen:
            uploaded += 1
            total_bytes += os.path.getsize(path)
            links.append("https://huggingface.co/%s/resolve/main/%s/%s"
                         % (repo, target, os.path.basename(path)))
    if uploaded:
        print("[заливка] готово: %d файлов, %.1f МБ" % (uploaded, total_bytes / 2**20), flush=True)
        for link in links[:6]:
            print("   " + link, flush=True)
        print("   каталог сборки: https://huggingface.co/%s/tree/main/builds/%s"
              % (repo, stamp), flush=True)
    return {"uploaded": uploaded, "bytes": total_bytes, "links": links,
            "folder": "builds/" + stamp}


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Kaggle PLUS")
    parser.add_argument("--demo", action="store_true", help="только самопроверка")
    parser.add_argument("--input-dir", help="папка с MIDI и аудио")
    parser.add_argument("--out-dir", default=os.path.join(WORK, "plus_out"))
    parser.add_argument("--stages", default="cleanup,quantize,hands,pedal,dynamics,report")
    parser.add_argument("--repo", default=REPO, help="репозиторий HF для сборки")
    parser.add_argument("--no-upload", action="store_true", help="не заливать сборку")
    parser.add_argument("--no-audio", action="store_true", help="не заливать wav превью")
    args = parser.parse_args(argv)

    ensure_deps()
    pack_dir = ensure_pack(WORK)
    report = demo(pack_dir)

    if args.input_dir and os.path.isdir(args.input_dir):
        command = [sys.executable, os.path.join(pack_dir, "plus_batch.py"),
                   "--input-dir", args.input_dir, "--out-dir", args.out_dir,
                   "--stages", args.stages]
        print("пачка: " + " ".join(command[1:]), flush=True)
        subprocess.run(command, check=False)
        report["batch_out"] = args.out_dir

    if args.no_upload:
        print("заливка отключена флагом --no-upload", flush=True)
        report["upload"] = {"uploaded": 0, "reason": "disabled"}
    else:
        stamp = __import__("datetime").datetime.utcnow().strftime("%Y%m%d-%H%M%S")
        folders = {"demo": os.path.join(WORK, "plus_demo")}
        if os.path.isdir(args.out_dir):
            folders["batch"] = args.out_dir
        report["upload"] = upload_build(folders, args.repo, stamp,
                                        with_audio=not args.no_audio)

    report_path = os.path.join(WORK, "plus_demo", "report.json")
    with open(report_path, "w", encoding="utf-8") as fh:
        json.dump(report, fh, ensure_ascii=False, indent=1, default=str)

    # отчёт дописывается после заливки, поэтому отправляем его ещё раз: в сборке
    # на HF должны быть видны и ссылки, и итоговые метрики
    uploaded = report.get("upload") or {}
    if uploaded.get("folder"):
        token = _hf_token()
        if token:
            from huggingface_hub import HfApi
            HfApi(token=token).upload_file(
                path_or_fileobj=report_path,
                path_in_repo=uploaded["folder"] + "/demo/report.json",
                repo_id=args.repo, commit_message="build report with upload links")
            print("[заливка] report.json обновлён в сборке", flush=True)
    return 0


if __name__ == "__main__":
    sys.exit(main())