File size: 2,572 Bytes
38bce11 6019d52 38bce11 6019d52 38bce11 6019d52 38bce11 6019d52 38bce11 6019d52 38bce11 6019d52 38bce11 6019d52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 | """Run the Kd-only frozen-encoder pass on a temporary Modal L4.
This script is intentionally not run by tests. Build the sample locally first,
review the current Modal price, then invoke it with:
modal run cloud/modal_embed.py
Retrieve the result with the `modal volume get` command printed on completion.
"""
from __future__ import annotations
import json
import os
import subprocess
from pathlib import Path
import modal
app = modal.App("mitointeract-v2-embeddings")
output_volume = modal.Volume.from_name(
"mitointeract-v2-artifacts", create_if_missing=True
)
cache_volume = modal.Volume.from_name("mitointeract-hf-cache", create_if_missing=True)
recovery_dir = Path(__file__).parents[1]
image = (
modal.Image.debian_slim(python_version="3.12")
.pip_install(
"numpy>=2.0",
"safetensors>=0.5",
"torch>=2.4",
"transformers>=4.45",
)
.add_local_dir(recovery_dir / "src", remote_path="/app/src")
.add_local_file(
recovery_dir / "scripts/embed_sample.py", remote_path="/app/embed_sample.py"
)
.add_local_file(
recovery_dir / "artifacts/balm-kd/sample.jsonl",
remote_path="/input/sample.jsonl",
)
)
@app.function(
image=image,
gpu="L4",
cpu=4,
memory=16_384,
timeout=3_600,
volumes={"/output": output_volume, "/cache": cache_volume},
)
def embed() -> dict:
env = os.environ.copy()
env["PYTHONPATH"] = "/app/src"
env["HF_HOME"] = "/cache/huggingface"
subprocess.run(
[
"python",
"/app/embed_sample.py",
"--sample",
"/input/sample.jsonl",
"--target-key",
"pkd",
"--target-name",
"pKd",
"--batch-size",
"32",
"--device",
"cuda",
"--output",
"/output/embeddings-balm-kd.npz",
],
check=True,
env=env,
)
output_volume.commit()
cache_volume.commit()
metadata = json.loads(Path("/output/embeddings-balm-kd.json").read_text())
metadata["output_bytes"] = Path("/output/embeddings-balm-kd.npz").stat().st_size
return metadata
@app.local_entrypoint()
def main() -> None:
result = embed.remote()
print(json.dumps(result, indent=2))
print(
"modal volume get mitointeract-v2-artifacts "
"embeddings-balm-kd.npz artifacts/embeddings-balm-kd.npz"
)
print(
"modal volume get mitointeract-v2-artifacts "
"embeddings-balm-kd.json artifacts/embeddings-balm-kd.json"
)
|