1
import json
2
import os
3
from nemo_platform import NeMoPlatform, ConflictError
4
5
NMP_BASE_URL = os.environ.get("NMP_BASE_URL", "http://localhost:8080")
6
client = NeMoPlatform(
7
    base_url=NMP_BASE_URL,
8
    workspace="default"
9
)

1
# Install required packages for dataset preparation
2
%pip install -q datasets huggingface_hub

1
import uuid
2
import numpy as np
3
4
# Demo query and documents for baseline comparison
5
DEMO_QUERY = "Conditional Random Fields: Probabilistic Models for Segmenting and Labeling Sequence Data"
6
7
DEMO_DOCS = [
8
    "Bidirectional LSTM-CRF Models for Sequence Tagging",   # CRF-based paper
9
    "An Introduction to Conditional Random Fields",         # CRF tutorial  
10
    "Random Forests",                                       # Keyword trap! Unrelated.
11
    "Neural Architectures for Named Entity Recognition",    # Related to sequence labeling; may use CRFs
12
    "Support Vector Machines for Classification",           # Unrelated ML method
13
]
14
15
DEMO_LABELS = ["BiLSTM-CRF", "CRF Tutorial", "Random Forest", "NER", "SVM"]
16
DEMO_RELEVANT = {0, 1, 3}  # Papers actually relevant to CRFs
17
18
def cosine_similarity(a, b):
19
    """Calculate cosine similarity between two vectors."""
20
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

1
# NGC API key is required to pull NIM images from nvcr.io
2
NGC_API_KEY = os.environ.get("NGC_API_KEY")
3
if not NGC_API_KEY:
4
    raise ValueError("NGC_API_KEY environment variable is required. Get one at https://ngc.nvidia.com/ → Setup → Generate API Key")
5
6
# Create NGC secret for pulling NIM images
7
NGC_SECRET_NAME = "ngc-api-key"
8
try:
9
    client.secrets.create(name=NGC_SECRET_NAME, workspace="default", value=NGC_API_KEY)
10
    print(f"Created secret: {NGC_SECRET_NAME}")
11
except ConflictError:
12
    print(f"Secret '{NGC_SECRET_NAME}' already exists, continuing...")
13
14
# Deploy base model for baseline comparison
15
BASE_MODEL_HF = "nvidia/llama-nemotron-embed-1b-v2"
16
NIM_IMAGE = "nvcr.io/nim/nvidia/llama-nemotron-embed-1b-v2"
17
NIM_TAG = "1.13.0"
18
19
baseline_suffix = uuid.uuid4().hex[:4]
20
BASELINE_DEPLOYMENT_CONFIG = f"baseline-embedding-cfg-{baseline_suffix}"
21
BASELINE_DEPLOYMENT_NAME = f"baseline-embedding-{baseline_suffix}"
22
23
print("Creating baseline deployment config...")
24
baseline_config = client.inference.deployment_configs.create(
25
    workspace="default",
26
    name=BASELINE_DEPLOYMENT_CONFIG,
27
    engine="nim",
28
    model_spec={},
29
    executor_config={
30
        "gpu": 1,
31
        "image_name": NIM_IMAGE,
32
        "image_tag": NIM_TAG,
33
    },
34
)
35
36
print("Deploying base model...")
37
baseline_deployment = client.inference.deployments.create(
38
    workspace="default",
39
    name=BASELINE_DEPLOYMENT_NAME,
40
    config=baseline_config.name
41
)
42
print(f"Baseline deployment: {baseline_deployment.name}")

1
import time
2
from IPython.display import clear_output
3
4
# Wait for baseline deployment
5
TIMEOUT_MINUTES = 15
6
start_time = time.time()
7
8
print(f"Waiting for baseline deployment...")
9
while True:
10
    status = client.inference.deployments.retrieve(
11
        name=BASELINE_DEPLOYMENT_NAME,
12
        workspace="default"
13
    )
14
    
15
    elapsed = time.time() - start_time
16
    elapsed_str = f"{int(elapsed//60)}m {int(elapsed%60)}s"
17
    
18
    clear_output(wait=True)
19
    print(f"Baseline deployment: {status.status} | {elapsed_str}")
20
    
21
    if status.status == "READY":
22
        print("Baseline model ready!")
23
        if not client.models.wait_for_gateway(BASELINE_DEPLOYMENT_NAME, workspace="default", timeout=60):
24
            raise RuntimeError("Inference gateway did not become ready")
25
        break
26
    if status.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
27
        raise RuntimeError(f"Baseline deployment failed: {status.status}")
28
    if elapsed > TIMEOUT_MINUTES * 60:
29
        raise TimeoutError("Baseline deployment timeout")
30
    
31
    time.sleep(10)

1
# Run baseline ranking with the base model
2
BASE_MODEL_ID = "nvidia/llama-nemotron-embed-1b-v2"
3
4
# Get query embedding
5
query_response = client.inference.gateway.provider.post(
6
    "v1/embeddings",
7
    name=BASELINE_DEPLOYMENT_NAME,
8
    workspace="default",
9
    body={
10
        "model": BASE_MODEL_ID,
11
        "input": [DEMO_QUERY],
12
        "input_type": "query"
13
    }
14
)
15
base_query_emb = query_response["data"][0]["embedding"]
16
17
# Get document embeddings
18
doc_response = client.inference.gateway.provider.post(
19
    "v1/embeddings",
20
    name=BASELINE_DEPLOYMENT_NAME,
21
    workspace="default",
22
    body={
23
        "model": BASE_MODEL_ID,
24
        "input": DEMO_DOCS,
25
        "input_type": "passage"
26
    }
27
)
28
base_doc_embs = [d["embedding"] for d in doc_response["data"]]
29
30
# Calculate similarities and rank
31
scores = [(i, cosine_similarity(base_query_emb, base_doc_embs[i])) for i in range(len(DEMO_DOCS))]
32
BASELINE_RANKING = sorted(scores, key=lambda x: -x[1])
33
34
# Display baseline results
35
print(f"Query: \"{DEMO_QUERY}\"\n")
36
print("Base Model Ranking:")
37
print("-" * 55)
38
for rank, (idx, score) in enumerate(BASELINE_RANKING, 1):
39
    marker = " <-- relevant" if idx in DEMO_RELEVANT else ""
40
    print(f"  #{rank}  [{score:.3f}]  {DEMO_LABELS[idx]}{marker}")

1
# Delete baseline deployment to free GPU for training
2
print("Deleting baseline deployment to free GPU...")
3
client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")
4
5
# Wait for deployment to be fully deleted before deleting config
6
print("Waiting for deployment deletion...")
7
while True:
8
    try:
9
        status = client.inference.deployments.retrieve(name=BASELINE_DEPLOYMENT_NAME, workspace="default")
10
        if status.status == "DELETED":
11
            break
12
        print(f"  Status: {status.status}")
13
        time.sleep(5)
14
    except Exception:
15
        # Deployment no longer exists
16
        break
17
18
# Now safe to delete the config
19
client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")
20
print("GPU freed. Proceed to fine-tune and improve these rankings.")

1
from pathlib import Path
2
from datasets import load_dataset
3
import json
4
5
# HuggingFace token for dataset access
6
HF_TOKEN = os.environ.get("HF_TOKEN")
7
if not HF_TOKEN:
8
    raise ValueError("HF_TOKEN environment variable is required. Get one at https://huggingface.co/settings/tokens")
9
os.environ["HF_TOKEN"] = HF_TOKEN
10
11
# Configuration
12
DATASET_SIZE = 3000        # Number of triplets (increase for better results, max ~684K)
13
VALIDATION_SPLIT = 0.05   # 5% held out for validation
14
SEED = 42
15
DATASET_PATH = Path("embedding-dataset").absolute()
16
17
# Create directory
18
os.makedirs(DATASET_PATH, exist_ok=True)
19
20
# Download SPECTER dataset
21
print("Downloading SPECTER dataset...")
22
data = load_dataset("embedding-data/SPECTER")["train"].shuffle(seed=SEED).select(range(DATASET_SIZE))
23
24
# Split into train/validation
25
print("Splitting into train/validation...")
26
splits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)
27
train_data = splits["train"]
28
validation_data = splits["test"]
29
30
# Convert to triplet JSONL format
31
print("Saving to JSONL...")
32
for name, dataset in [("training", train_data), ("validation", validation_data)]:
33
    with open(f"{DATASET_PATH}/{name}.jsonl", "w") as f:
34
        for row in dataset:
35
            # SPECTER format: row['set'] = [query, positive, negative]
36
            triplet = {
37
                "query": row["set"][0],
38
                "pos_doc": row["set"][1],
39
                "neg_doc": [row["set"][2]]  # List of negative documents
40
            }
41
            f.write(json.dumps(triplet) + "\n")
42
43
print(f"\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples")
44
print(f"\nExample triplet:")
45
print(f"  Query:    {train_data[0]['set'][0][:100]}...")
46
print(f"  Positive: {train_data[0]['set'][1][:100]}...")
47
print(f"  Negative: {train_data[0]['set'][2][:100]}...")

1
# Create fileset to store embedding training data
2
DATASET_NAME = "embedding-dataset"
3
4
try:
5
    client.files.filesets.create(
6
        workspace="default",
7
        name=DATASET_NAME,
8
        description="SPECTER embedding training data (scientific paper triplets)"
9
    )
10
    print(f"Created fileset: {DATASET_NAME}")
11
except ConflictError:
12
    print(f"Fileset '{DATASET_NAME}' already exists, continuing...")
13
14
# Upload training data files
15
client.files.upload(
16
    local_path=f"{DATASET_PATH}/",
17
    remote_path="",
18
    fileset=DATASET_NAME,
19
    workspace="default"
20
)
21
22
# Validate upload
23
print("\nUploaded files:")
24
print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace="default").data], indent=2))

1
# Create secrets for model access
2
# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)
3
HF_TOKEN = os.getenv("HF_TOKEN")
4
5
6
def create_or_get_secret(name: str, value: str | None, label: str):
7
    if not value:
8
        raise ValueError(f"{label} is not set")
9
    try:
10
        secret = client.secrets.create(
11
            name=name,
12
            workspace="default",
13
            value=value,
14
        )
15
        print(f"Created secret: {name}")
16
        return secret
17
    except ConflictError:
18
        print(f"Secret '{name}' already exists, continuing...")
19
        return client.secrets.retrieve(name=name, workspace="default")
20
21
22
# Create HuggingFace token secret (for downloading model from HF during training)
23
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN")
24
print(f"HF_TOKEN secret: {hf_secret.name}")
25
26
# NGC secret was already created in baseline step (Step 2), or use the platform default
27
if "NGC_SECRET_NAME" not in globals():
28
    NGC_SECRET_NAME = "ngc-api-key"
29
print(f"NGC_API_KEY secret: {NGC_SECRET_NAME}")

1
import time
2
from nemo_platform.types.files import HuggingfaceStorageConfigParam
3
4
HF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"
5
MODEL_NAME = "nv-nemotron-embed-1b-base"
6
7
# Ensure you have a HuggingFace token secret created
8
try:
9
    base_model_fs = client.files.filesets.create(
10
        workspace="default",
11
        name=MODEL_NAME,
12
        description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",
13
        storage=HuggingfaceStorageConfigParam(
14
            type="huggingface",
15
            # repo_id is the full model name from Hugging Face
16
            repo_id=HF_REPO_ID,
17
            repo_type="model",
18
            # we use the secret created in the previous step
19
            token_secret=hf_secret.name
20
        )
21
    )
22
except ConflictError as e:
23
    print(f"Base model fileset already exists. Skipping creation.")
24
    base_model_fs = client.files.filesets.retrieve(
25
        workspace="default",
26
        name=MODEL_NAME,
27
    )
28
29
# Create Model Entity referencing the FileSet
30
try:
31
    base_model = client.models.create(
32
        workspace="default",
33
        name=MODEL_NAME,
34
        fileset=f"default/{MODEL_NAME}",
35
        trust_remote_code=True,
36
    )
37
    print(f"Created Model Entity: {MODEL_NAME}")
38
except ConflictError:
39
    print(f"Base model already exists. Updating fileset if different.")
40
    base_model = client.models.update(
41
        workspace="default",
42
        name=MODEL_NAME,
43
        fileset=f"default/{MODEL_NAME}",
44
        trust_remote_code=True,
45
    )
46
47
print(f"\nBase model fileset: fileset://default/{base_model.name}")
48
print("\nBase model files:")
49
print(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))
50
51
# Wait for ModelSpec to be populated from the checkpoint
52
print("\nWaiting for ModelSpec to be populated...")
53
SPEC_TIMEOUT_SECONDS = 120
54
spec_start = time.time()
55
while not base_model.spec:
56
    if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:
57
        raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")
58
    time.sleep(2)
59
    base_model = client.models.retrieve(
60
        workspace="default",
61
        name=MODEL_NAME,
62
    )
63
64
print(f"ModelSpec populated: {base_model.spec}")

1
training={
2
    "training_type": "sft",
3
    "finetuning_type": "lora_merged",
4
    "lora": {"rank": 16, "alpha": 32},
5
    "max_seq_length": MAX_SEQ_LENGTH,
6
}

1
import uuid
2
from nemo_automodel_plugin.schema import AutomodelJobInput
3
4
job_suffix = uuid.uuid4().hex[:4]
5
JOB_NAME = f"embedding-finetune-job-{job_suffix}"
6
OUTPUT_NAME = f"nv-embed-finetuned-{job_suffix}"
7
8
EPOCHS = 1
9
BATCH_SIZE = 128
10
LEARNING_RATE = 5e-6
11
MAX_SEQ_LENGTH = 512
12
13
spec = AutomodelJobInput(
14
    model=f"default/{base_model.name}",
15
    dataset={"training": f"default/{DATASET_NAME}"},
16
    training={
17
        "training_type": "sft",
18
        "finetuning_type": "all_weights",
19
        "max_seq_length": MAX_SEQ_LENGTH,
20
    },
21
    schedule={"epochs": EPOCHS},
22
    batch={"global_batch_size": BATCH_SIZE, "micro_batch_size": 1},
23
    optimizer={"learning_rate": LEARNING_RATE},
24
    parallelism={"num_gpus_per_node": 1},
25
    output={"name": OUTPUT_NAME},
26
)
27
28
job = client.customization.automodel.jobs.create(
29
    spec=spec, workspace="default", name=JOB_NAME
30
)
31
32
print(f"Submitted job: {job.job.name}")
33
print(f"Output model: {OUTPUT_NAME}")

1
import time
2
from IPython.display import clear_output
3
4
# Poll job status every 10 seconds until completed
5
while True:
6
    status = client.jobs.get_status(
7
        name=job.job.name,
8
        workspace="default"
9
    )
10
    
11
    clear_output(wait=True)
12
    print(f"Job Status: {status.model_dump_json(indent=2)}")
13
14
    # Extract training progress from nested steps structure
15
    step: int | None = None
16
    max_steps: int | None = None
17
    training_phase: str | None = None
18
19
    for job_step in status.steps or []:
20
        if job_step.name == "training":
21
            for task in job_step.tasks or []:
22
                task_details = task.status_details or {}
23
                step = task_details.get("step")
24
                max_steps = task_details.get("max_steps")
25
                training_phase = task_details.get("phase")
26
                break
27
            break
28
29
    if step is not None and max_steps is not None:
30
        progress_pct = (step / max_steps) * 100
31
        print(f"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)")
32
        if training_phase:
33
            print(f"Training Phase: {training_phase}")
34
    else:
35
        print("Training step not started yet or progress info not available")
36
    
37
    # Exit loop when job is completed (or failed/cancelled)
38
    if status.status in ("completed", "failed", "cancelled", "error"):
39
        print(f"\nJob finished with status: {status.status}")
40
        break
41
    
42
    time.sleep(10)

1
# Validate model entity exists
2
model_entity = client.models.retrieve(workspace="default", name=OUTPUT_NAME)
3
print(model_entity.model_dump_json(indent=2))

1
# Create deployment config for embedding model
2
deploy_suffix = uuid.uuid4().hex[:4]
3
DEPLOYMENT_CONFIG_NAME = f"embedding-model-deployment-cfg-{deploy_suffix}"
4
DEPLOYMENT_NAME = f"embedding-model-deployment-{deploy_suffix}"
5
6
# Embedding NIM image
7
NIM_IMAGE = "nvcr.io/nim/nvidia/llama-nemotron-embed-1b-v2"
8
NIM_TAG = "1.13.0"  # Update if using newer NIM release
9
10
deployment_config = client.inference.deployment_configs.create(
11
    workspace="default",
12
    name=DEPLOYMENT_CONFIG_NAME,
13
    engine="nim",
14
    model_spec={
15
        "model_namespace": "default",
16
        "model_name": OUTPUT_NAME,
17
    },
18
    executor_config={
19
        "gpu": 1,
20
        "image_name": NIM_IMAGE,
21
        "image_tag": NIM_TAG,
22
    },
23
)
24
25
# Deploy model
26
deployment = client.inference.deployments.create(
27
    workspace="default",
28
    name=DEPLOYMENT_NAME,
29
    config=deployment_config.name
30
)
31
32
print(f"Deployment name: {deployment.name}")
33
print(f"Deployment status: {client.inference.deployments.retrieve(name=deployment.name, workspace='default').status}")

1
import time
2
from IPython.display import clear_output
3
4
# Poll deployment status every 15 seconds until ready
5
TIMEOUT_MINUTES = 30
6
start_time = time.time()
7
timeout_seconds = TIMEOUT_MINUTES * 60
8
9
print(f"Monitoring deployment '{deployment.name}'...")
10
print(f"Timeout: {TIMEOUT_MINUTES} minutes\n")
11
12
while True:
13
    deployment_status = client.inference.deployments.retrieve(
14
        name=deployment.name,
15
        workspace="default"
16
    )
17
    
18
    elapsed = time.time() - start_time
19
    elapsed_min = int(elapsed // 60)
20
    elapsed_sec = int(elapsed % 60)
21
    
22
    clear_output(wait=True)
23
    print(f"Deployment: {deployment.name}")
24
    print(f"Status: {deployment_status.status}")
25
    print(f"Elapsed time: {elapsed_min}m {elapsed_sec}s")
26
    
27
    # Check if deployment is ready
28
    if deployment_status.status == "READY":
29
        print("\nDeployment is ready!")
30
        if not client.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):
31
            raise RuntimeError("Inference gateway did not become ready")
32
        break
33
    
34
    # Check for failure states
35
    if deployment_status.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
36
        raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")
37
    
38
    # Check timeout
39
    if elapsed > timeout_seconds:
40
        raise TimeoutError(f"Deployment timeout after {TIMEOUT_MINUTES} minutes")
41
    
42
    time.sleep(15)

1
# Compare: same query, base model vs fine-tuned
2
# Using the same DEMO_QUERY and DEMO_DOCS from the baseline test
3
MODEL_ID = f"default/{OUTPUT_NAME}"
4
5
# Get query embedding from fine-tuned model
6
query_response = client.inference.gateway.provider.post(
7
    "v1/embeddings",
8
    name=deployment.name,
9
    workspace="default",
10
    body={
11
        "model": MODEL_ID,
12
        "input": [DEMO_QUERY],
13
        "input_type": "query"
14
    }
15
)
16
query_embedding = query_response["data"][0]["embedding"]
17
18
# Get document embeddings from fine-tuned model
19
doc_response = client.inference.gateway.provider.post(
20
    "v1/embeddings",
21
    name=deployment.name,
22
    workspace="default",
23
    body={
24
        "model": MODEL_ID,
25
        "input": DEMO_DOCS,
26
        "input_type": "passage"
27
    }
28
)
29
doc_embeddings = [d["embedding"] for d in doc_response["data"]]
30
31
# Calculate similarities and rank
32
scores = [(i, cosine_similarity(query_embedding, doc_embeddings[i])) for i in range(len(DEMO_DOCS))]
33
FINETUNED_RANKING = sorted(scores, key=lambda x: -x[1])
34
35
# Display side-by-side comparison
36
print(f"Query: \"{DEMO_QUERY}\"\n")
37
print(f"{'Rank':<6} {'Base Model':<30} {'Fine-tuned Model':<30}")
38
print("-" * 66)
39
40
for rank in range(len(DEMO_DOCS)):
41
    b_idx, b_score = BASELINE_RANKING[rank]
42
    f_idx, f_score = FINETUNED_RANKING[rank]
43
    
44
    b_label = f"{DEMO_LABELS[b_idx]} [{b_score:.3f}]" + (" *" if b_idx in DEMO_RELEVANT else "")
45
    f_label = f"{DEMO_LABELS[f_idx]} [{f_score:.3f}]" + (" *" if f_idx in DEMO_RELEVANT else "")
46
    
47
    print(f"#{rank+1:<5} {b_label:<30} {f_label:<30}")
48
49
print("\n* = relevant paper")
50
print("\nThe fine-tuned model pushes 'Random Forest' down and ranks CRF papers higher.")