# パッケージ更新。Jupyterでは pip install をそのまま書かず、Python経由で実行します。
import sys
import subprocess

subprocess.check_call([
    sys.executable,
    "-m",
    "pip",
    "install",
    "-U",
    "litert-lm",
    "litert-lm-api",
    "huggingface_hub",
    "psutil",
])

from pathlib import Path
from huggingface_hub import hf_hub_download
import os
import litert_lm
import psutil

MODEL_REPO = "litert-community/gemma-4-E2B-it-litert-lm"
MODEL_FILE = "gemma-4-E2B-it.litertlm"
MODEL_DIR = Path("models") / "gemma-4-E2B-it-litert-lm"
CACHE_DIR = Path("litert_cache")

# snapshot_download だとリポジトリ全体(約22GB)を落としに行くため、必要な1ファイルだけ取得します。
model_path = Path(hf_hub_download(
    repo_id=MODEL_REPO,
    filename=MODEL_FILE,
    local_dir=MODEL_DIR,
))

litert_lm.set_min_log_severity(litert_lm.LogSeverity.INFO)

def format_bytes(size):
    for unit in ["B", "KB", "MB", "GB", "TB"]:
        if size < 1024:
            return f"{size:.1f} {unit}"
        size /= 1024
    return f"{size:.1f} PB"

def process_ram_mb():
    return psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024

def print_memory(label):
    print(f"[{label}] Python RAM: {process_ram_mb():.1f} MiB")

print(f"Model file: {model_path}")
print(f"Model file size: {format_bytes(model_path.stat().st_size)}")
print_memory("initial")

def run_cpu_once():
    print("\n=== CPU backend ===")
    ram_before = process_ram_mb()
    print_memory("CPU before engine")
    with litert_lm.Engine(
        str(model_path),
        backend=litert_lm.Backend.CPU(),
        cache_dir=str(CACHE_DIR / "cpu"),
        enable_speculative_decoding=False,
    ) as engine:
        print_memory("CPU after engine")
        with engine.create_conversation(
            messages=[litert_lm.Message.system("あなたは日本語で簡潔に答えるアシスタントです。")]
        ) as chat:
            response = chat.send_message("日本語で自己紹介して")
            print_memory("CPU after response")
            print(response["content"][0]["text"])
    print_memory("CPU after close")
    print(f"CPU Python RAM delta: {process_ram_mb() - ram_before:+.1f} MiB")

run_cpu_once()