Embedding Model Customization

import json
import os
import time
import uuid
from IPython.display import clear_output
from nemo_platform import NeMoPlatform, ConflictError
NMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080")
client = NeMoPlatform(base_url=NMP_BASE_URL, workspace="default")
print(f"Platform: {NMP_BASE_URL}")

from nemo_platform.types.files import HuggingfaceStorageConfigParam
NGC_API_KEY = os.environ.get("NGC_API_KEY")
if not NGC_API_KEY:
    raise ValueError("NGC_API_KEY is required to pull Retriever NIM images from nvcr.io")
NGC_SECRET_NAME = "ngc-api-key"
try:
    client.secrets.create(name=NGC_SECRET_NAME, workspace="default", value=NGC_API_KEY)
    print(f"Created secret: {NGC_SECRET_NAME}")
except ConflictError:
    print(f"Secret '{NGC_SECRET_NAME}' already exists")
HF_TOKEN = os.getenv("HF_TOKEN")
HF_REPO_ID = "nvidia/Nemotron-3-Embed-1B-BF16"
MODEL_NAME = "nemotron-3-embed-1b"
storage_kwargs = {"type": "huggingface", "repo_id": HF_REPO_ID, "repo_type": "model"}
if HF_TOKEN:
    try:
        client.secrets.create(name="hf-token", workspace="default", value=HF_TOKEN)
    except ConflictError:
        pass
    storage_kwargs["token_secret"] = "hf-token"
storage = HuggingfaceStorageConfigParam(**storage_kwargs)
try:
    client.files.filesets.create(
        workspace="default",
        name=MODEL_NAME,
        purpose="model",
        description="Nemotron 3 Embed 1B BF16 checkpoint",
        storage=storage,
    )
except ConflictError:
    print("Base model fileset already exists")
try:
    base_model = client.models.create(
        workspace="default",
        name=MODEL_NAME,
        fileset=f"default/{MODEL_NAME}",
        trust_remote_code=True,
    )
    print(f"Created Model Entity: {MODEL_NAME}")
except ConflictError:
    base_model = client.models.update(
        workspace="default",
        name=MODEL_NAME,
        fileset=f"default/{MODEL_NAME}",
        trust_remote_code=True,
    )
    print(f"Updated Model Entity: {MODEL_NAME}")
print("Waiting for ModelSpec...")
spec_start = time.time()
while not getattr(base_model, "spec", None):
    if time.time() - spec_start > 180:
        raise TimeoutError("ModelSpec not populated")
    time.sleep(3)
    base_model = client.models.retrieve(workspace="default", name=MODEL_NAME)
print(f"fileset={base_model.fileset} spec={base_model.spec}")
if not base_model.fileset:
    raise RuntimeError("Entity has a null fileset; it cannot be trained. Recreate from the Hugging Face fileset.")

from nemo_data_designer_plugin.jobs.retrieval_spec import RetrievalPrepareJobConfig
NVDOCS_SDG = "hf://nvidia/Retrieval-Synthetic-NVDocs-v1@1c0d1856f3fb595b2dda98d4b61061fa6d782d51/nv_pp_dd_sdg.json"
prepare_spec = RetrievalPrepareJobConfig(
    sdg_input=NVDOCS_SDG,
    enable_mining=True,
    model=f"default/{MODEL_NAME}",
    hard_negatives_to_mine=5,
    hard_neg_margin=0.95,
    mining={
        "query_embedding_batch_size": 64,
        "document_embedding_batch_size": 64,
    },
)
prepare_job = client.data_designer.retrieval_prepare(prepare_spec, workspace="default")
print(f"Prepare job: {prepare_job.name}")
prepare_job.wait_until_done()
prepare_status = prepare_job.get_job_status()
print(f"Stage 1 status: {prepare_status}")
if prepare_status != "completed":
    raise RuntimeError(f"retrieval-prepare finished with status {prepare_status}")

import csv
import tempfile
from pathlib import Path
results = client.jobs.results.list(prepare_job.name, workspace="default").data
artifacts = next(row for row in results if row.name == "artifacts")
ARTIFACTS_REF = artifacts.artifact_url
print("Stage 1 artifacts:", ARTIFACTS_REF)
TRAINING_FILESET = "nvdocs-embed-artifacts"
fileset_ref, _, remote_prefix = ARTIFACTS_REF.partition("#")
job_fileset = fileset_ref.split("/", 1)[-1]
local_dir = Path(tempfile.mkdtemp(prefix="nvdocs-artifacts-"))
client.files.download(
    remote_path=f"{remote_prefix}/training.jsonl",
    local_path=str(local_dir / "training.jsonl"),
    fileset=job_fileset,
    workspace="default",
)
# eval_beir arrives under local_dir; a destination already named eval_beir nests as eval_beir/eval_beir
client.files.download(
    remote_path=f"{remote_prefix}/eval_beir",
    local_path=str(local_dir),
    fileset=job_fileset,
    workspace="default",
)
if not (local_dir / "eval_beir" / "corpus.jsonl").is_file():
    raise RuntimeError(f"eval_beir/corpus.jsonl missing under {local_dir}")
# the BEIR loader rejects duplicate judgments
qrels_path = local_dir / "eval_beir" / "qrels" / "test.tsv"
with qrels_path.open(encoding="utf-8", newline="") as stream:
    reader = csv.DictReader(stream, delimiter="\t")
    header = reader.fieldnames
    seen: set[tuple[str, str]] = set()
    kept: list[dict[str, str]] = []
    for row in reader:
        key = (row["query-id"], row["corpus-id"])
        if key not in seen:
            seen.add(key)
            kept.append(row)
with qrels_path.open("w", encoding="utf-8", newline="") as stream:
    writer = csv.DictWriter(stream, fieldnames=header, delimiter="\t", lineterminator="\n")
    writer.writeheader()
    writer.writerows(kept)
print(f"qrels judgments: {len(kept)}")
try:
    client.files.filesets.create(
        workspace="default",
        name=TRAINING_FILESET,
        purpose="dataset",
        description="NVDocs mined retrieval artifacts",
    )
except ConflictError:
    pass
client.files.upload(
    local_path=f"{local_dir}/",
    remote_path="",
    fileset=TRAINING_FILESET,
    workspace="default",
)
TRAINING_REF = f"default/{TRAINING_FILESET}"
print("dataset.training and retrieve-eval dataset:", TRAINING_REF)

from nemo_automodel_plugin.schema import AutomodelJobInput
job_suffix = uuid.uuid4().hex[:4]
JOB_NAME = f"embed-nvdocs-{job_suffix}"
OUTPUT_NAME = f"nemotron-3-embed-1b-tuned-{job_suffix}"
spec = AutomodelJobInput(
    model=f"default/{MODEL_NAME}",
    dataset={"training": TRAINING_REF},
    training={
        "training_type": "sft",
        "recipe": "bi_encoder",
        "finetuning_type": "all_weights",
        "max_seq_length": 512,
        "retrieval": {
            "query_max_length": 512,
            "passage_max_length": 512,
            "train_n_passages": 5,
            "export": {"primary": "hf"},
        },
    },
    schedule={"epochs": 1},
    batch={"global_batch_size": 128, "micro_batch_size": 4},
    optimizer={"learning_rate": 1e-5, "warmup_steps": 5, "weight_decay": 0.01},
    parallelism={"num_gpus_per_node": 1},
    output={"name": OUTPUT_NAME},
)
job = client.customization.automodel.jobs.create(spec=spec, workspace="default", name=JOB_NAME)
print(f"Submitted {job.job.name}")
print(f"Output model: {OUTPUT_NAME}")

while True:
    status = client.jobs.get_status(name=job.job.name, workspace="default")
    clear_output(wait=True)
    print(status.model_dump_json(indent=2))
    if status.status in ("completed", "failed", "cancelled", "error"):
        break
    time.sleep(30)
if status.status != "completed":
    raise RuntimeError(f"Training finished with status: {status.status}")

NIM_IMAGE = "nvcr.io/nim/nvidia/nemotron-3-embed-1b"
NIM_TAG = "2.2.0"
deploy_suffix = uuid.uuid4().hex[:4]
BASE_CFG = f"embed-base-cfg-{deploy_suffix}"
TUNED_CFG = f"embed-tuned-cfg-{deploy_suffix}"
BASE_DEP = f"embed-base-{deploy_suffix}"
TUNED_DEP = f"embed-tuned-{deploy_suffix}"
executor = {
    "gpu": 1,
    "image_name": NIM_IMAGE,
    "image_tag": NIM_TAG,
    "override_config": {"nimLegacy": False},
}
base_cfg = client.inference.deployment_configs.create(
    workspace="default",
    name=BASE_CFG,
    engine="nim",
    model_spec={},
    executor_config=executor,
    model_entity_id=f"default/{MODEL_NAME}",
)
tuned_cfg = client.inference.deployment_configs.create(
    workspace="default",
    name=TUNED_CFG,
    engine="nim",
    model_spec={"model_namespace": "default", "model_name": OUTPUT_NAME},
    executor_config=executor,
    model_entity_id=f"default/{OUTPUT_NAME}",
)
base_deployment = client.inference.deployments.create(
    workspace="default", name=BASE_DEP, config=base_cfg.name
)
tuned_deployment = client.inference.deployments.create(
    workspace="default", name=TUNED_DEP, config=tuned_cfg.name
)
print(base_deployment.name, tuned_deployment.name)

for deployment in (BASE_DEP, TUNED_DEP):
    if not client.models.wait_for_status(deployment, "READY", workspace="default", timeout=2700):
        raise RuntimeError(f"{deployment} did not reach READY")
for entity in (MODEL_NAME, OUTPUT_NAME):
    providers = client.models.retrieve(workspace="default", name=entity).model_providers
    print(entity, "providers:", providers)
    if not providers:
        raise RuntimeError(f"{entity} has no model_providers; retrieve-eval cannot score it")

resp = client.inference.gateway.provider.post(
    "v1/embeddings",
    name=TUNED_DEP,
    workspace="default",
    body={
        "model": f"default/{OUTPUT_NAME}",
        "input": ["What does NVLink connect?"],
        "input_type": "query",
    },
)
print("dim", len(resp["data"][0]["embedding"]))

import shutil
import subprocess
eval_spec = {
    "dataset": TRAINING_REF,
    "target": f"default/{OUTPUT_NAME}",
    "baseline": f"default/{MODEL_NAME}",
    "k": [1, 5, 10, 100],
}
cmd = [
    shutil.which("nemo") or "nemo",
    "evaluator",
    "retrieve-eval",
    "submit",
    "--workspace",
    "default",
    "--spec",
    json.dumps(eval_spec),
]
proc = subprocess.run(cmd, check=False, capture_output=True, text=True)
if proc.returncode != 0:
    raise RuntimeError(f"retrieve-eval submit failed:\n{proc.stderr}")
eval_job_name = json.loads(proc.stdout[proc.stdout.index("{") :])["name"]
print("retrieve-eval job:", eval_job_name)
while True:
    status = client.jobs.get_status(name=eval_job_name, workspace="default")
    print("eval status:", status.status)
    if status.status in ("completed", "failed", "cancelled", "error"):
        break
    time.sleep(30)
if status.status != "completed":
    raise RuntimeError(f"retrieve-eval finished with status: {status.status}")
for row in client.jobs.results.list(eval_job_name, workspace="default").data:
    print(json.dumps(row.model_dump(), default=str, indent=2)[:4000])