Embedding Model Customization

import os
os.environ["NHX_BASE_URL"] = "<YOUR_NHX_BASE_URL>"

import json
import os
import time
import uuid
from IPython.display import clear_output
from nemo_helix_plugin.client.client import NemoClient
from nemo_helix_plugin.client.errors import ConflictError
from nemo_helix_plugin.files.storage_config import HuggingfaceStorageConfig
from nemo_helix_plugin.files.types import CreateFilesetRequest, FilesetPurpose
from nemo_helix_plugin.models.types import (
    ContainerExecutorConfig,
    CreateModelDeploymentConfigRequest,
    CreateModelDeploymentRequest,
    CreateModelEntityRequest,
    Engine,
    ModelDeploymentConfigModelSpec,
    UpdateModelEntityRequest,
)
from nemo_helix_plugin.secrets.types import HelixSecretCreateRequest
from nemo_customizer.sdk import Customization
from nemo_data_designer_plugin.sdk.resources import DataDesignerResource
from filesets import FilesetFileSystem
from nemo_helix_plugin.files.client import FilesClient
from pydantic import SecretStr
from nemo_helix_plugin.schema import SecretRef
client = NemoClient.from_config()
customization = Customization.from_client(client)
data_designer = DataDesignerResource(client)
print(f"Platform: {client.base_url}")

from nemo_helix_plugin.files.storage_config import HuggingfaceStorageConfig
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_secret(body=HelixSecretCreateRequest(name=NGC_SECRET_NAME, value=SecretStr(NGC_API_KEY))).data()
    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"
if HF_TOKEN:
    try:
        client.secrets.create_secret(body=HelixSecretCreateRequest(name="hf-token", value=SecretStr(HF_TOKEN))).data()
    except ConflictError:
        pass
    storage = HuggingfaceStorageConfig(
        repo_id=HF_REPO_ID,
        repo_type="model",
        token_secret=SecretRef("hf-token"),
    )
else:
    storage = HuggingfaceStorageConfig(
        repo_id=HF_REPO_ID,
        repo_type="model",
    )
try:
    client.files.create_fileset(
        body=CreateFilesetRequest(
            name=MODEL_NAME,
            purpose=FilesetPurpose.MODEL,
            description="Nemotron 3 Embed 1B BF16 checkpoint",
            storage=storage,
        ),
    ).data()
except ConflictError:
    print("Base model fileset already exists")
try:
    base_model = client.models.create_model(
        body=CreateModelEntityRequest(
            name=MODEL_NAME,
            fileset=f"default/{MODEL_NAME}",
            trust_remote_code=True,
        ),
    ).data()
    print(f"Created Model Entity: {MODEL_NAME}")
except ConflictError:
    base_model = client.models.update_model(
        name=MODEL_NAME,
        body=UpdateModelEntityRequest(
            fileset=f"default/{MODEL_NAME}",
            trust_remote_code=True,
        ),
    ).data()
    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.get_model(name=MODEL_NAME).data()
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 RetrievalMiningOptions, 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=RetrievalMiningOptions(
        query_embedding_batch_size=64,
        document_embedding_batch_size=64,
    ),
)
prepare_job = 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-"))
fs = FilesetFileSystem(client=FilesClient.from_client(client))
fs.get(
    rpath=f"default/{job_fileset}/{remote_prefix}/training.jsonl",
    lpath=str(local_dir / "training.jsonl"),
)
# eval_beir arrives under local_dir; a destination already named eval_beir nests as eval_beir/eval_beir
fs.get(
    rpath=f"default/{job_fileset}/{remote_prefix}/eval_beir",
    lpath=str(local_dir),
    recursive=True,
)
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
    if header is None:
        raise RuntimeError(f"{qrels_path} is missing a header")
    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.create_fileset(
        body=CreateFilesetRequest(
            name=TRAINING_FILESET,
            purpose=FilesetPurpose.DATASET,
            description="NVDocs mined retrieval artifacts",
        ),
    ).data()
except ConflictError:
    pass
fs.put(
    lpath=f"{local_dir}/",
    rpath=f"default/{TRAINING_FILESET}/",
    recursive=True,
)
TRAINING_REF = f"default/{TRAINING_FILESET}"
print("dataset.training and retrieve-eval dataset:", TRAINING_REF)

from nemo_automodel_plugin.schema import (
    AutomodelJobInput,
    BatchSpec,
    DatasetSpec,
    ExportSpec,
    OptimizerSpec,
    OutputRequest,
    ParallelismSpec,
    RetrievalSpec,
    ScheduleSpec,
    TrainingSpec,
)
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=DatasetSpec(training=TRAINING_REF),
    training=TrainingSpec(
        training_type="sft",
        recipe="bi_encoder",
        finetuning_type="all_weights",
        execution_profile="gpu",
        max_seq_length=512,
        precision="bf16",
        attn_implementation="flash_attention_2",
        retrieval=RetrievalSpec(
            query_max_length=512,
            passage_max_length=512,
            train_n_passages=5,
            do_distributed_inbatch_negative=False,
            export=ExportSpec(primary="hf"),
        ),
    ),
    schedule=ScheduleSpec(
        epochs=3,
        seed=42,
        validation_split=0.01,
        val_check_interval=0.25,
        checkpoint_selection="both",
    ),
    batch=BatchSpec(global_batch_size=128, micro_batch_size=8),
    optimizer=OptimizerSpec(
        learning_rate=2.5e-5,
        min_learning_rate=1.25e-5,
        warmup_steps=432,
        weight_decay=0.01,
        lr_decay_style="cosine",
    ),
    parallelism=ParallelismSpec(num_gpus_per_node=4, tensor_parallel_size=1),
    output=OutputRequest(name=OUTPUT_NAME),
)
job = customization.automodel.jobs.create(
    spec=spec.model_dump(mode="json", exclude_unset=True),
    workspace="default",
    name=JOB_NAME,
)
print(f"Submitted job: {JOB_NAME}")
print(f"Output model: {OUTPUT_NAME}")

import time
# Poll the job until it reaches a terminal state
while True:
    status = client.jobs.get_status(name=JOB_NAME, workspace="default")
    print(f"Status: {status.status}")
    if status.status.lower() in ("completed", "failed", "cancelled", "error"):
        break
    time.sleep(30)

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 = ContainerExecutorConfig(
    gpu=1,
    image_name=NIM_IMAGE,
    image_tag=NIM_TAG,
    override_config={"nimLegacy": False},
)
base_cfg = client.models.create_deployment_config(
    body=CreateModelDeploymentConfigRequest(
        name=BASE_CFG,
        engine=Engine.NIM,
        model_spec=ModelDeploymentConfigModelSpec(),
        executor_config=executor,
        model_entity_id=f"default/{MODEL_NAME}",
    ),
).data()
tuned_cfg = client.models.create_deployment_config(
    body=CreateModelDeploymentConfigRequest(
        name=TUNED_CFG,
        engine=Engine.NIM,
        model_spec=ModelDeploymentConfigModelSpec(
            model_namespace="default",
            model_name=OUTPUT_NAME,
        ),
        executor_config=executor,
        model_entity_id=f"default/{OUTPUT_NAME}",
    ),
).data()
base_deployment = client.models.create_deployment(
    body=CreateModelDeploymentRequest(name=BASE_DEP, config=base_cfg.name),
).data()
tuned_deployment = client.models.create_deployment(
    body=CreateModelDeploymentRequest(name=TUNED_DEP, config=tuned_cfg.name),
).data()
print(base_deployment.name, tuned_deployment.name)

for deployment in (BASE_DEP, TUNED_DEP):
    if not client.models.wait_for_deployment_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.get_model(name=entity).data().model_providers
    print(entity, "providers:", providers)
    if not providers:
        raise RuntimeError(f"{entity} has no model_providers; retrieve-eval cannot score it")

from typing import Any, cast
from nemo_helix_plugin.inference_gateway.client import InferenceGatewayClient
from nemo_helix_plugin.inference_gateway.types import JsonBody
resp = InferenceGatewayClient.from_client(client).provider_post(
    name=TUNED_DEP,
    workspace="default",
    trailing_uri="v1/embeddings",
    body=JsonBody(
        {
            "model": f"default/{OUTPUT_NAME}",
            "input": ["What does NVLink connect?"],
            "input_type": "query",
        }
    ),
).data()
resp_body = cast(dict[str, Any], resp)
print("dim", len(resp_body["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",
    "--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 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])