Embedding Model Customization

import json
import os
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"
)

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

# NGC API key is required to pull NIM images from nvcr.io
NGC_API_KEY = os.environ.get("NGC_API_KEY")
if not NGC_API_KEY:
    raise ValueError("NGC_API_KEY environment variable is required. Get one at https://ngc.nvidia.com/ → Setup → Generate API Key")
# Create NGC secret for pulling NIM images
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, continuing...")
# Deploy base model for baseline comparison
BASE_MODEL_HF = "nvidia/llama-nemotron-embed-1b-v2"
NIM_IMAGE = "nvcr.io/nim/nvidia/llama-nemotron-embed-1b-v2"
NIM_TAG = "1.13.0"
baseline_suffix = uuid.uuid4().hex[:4]
BASELINE_DEPLOYMENT_CONFIG = f"baseline-embedding-cfg-{baseline_suffix}"
BASELINE_DEPLOYMENT_NAME = f"baseline-embedding-{baseline_suffix}"
print("Creating baseline deployment config...")
baseline_config = client.inference.deployment_configs.create(
    workspace="default",
    name=BASELINE_DEPLOYMENT_CONFIG,
    engine="nim",
    model_spec={},
    executor_config={
        "gpu": 1,
        "image_name": NIM_IMAGE,
        "image_tag": NIM_TAG,
    },
)
print("Deploying base model...")
baseline_deployment = client.inference.deployments.create(
    workspace="default",
    name=BASELINE_DEPLOYMENT_NAME,
    config=baseline_config.name
)
print(f"Baseline deployment: {baseline_deployment.name}")

import time
from IPython.display import clear_output
# Wait for baseline deployment
TIMEOUT_MINUTES = 15
start_time = time.time()
print(f"Waiting for baseline deployment...")
while True:
    status = client.inference.deployments.retrieve(
        name=BASELINE_DEPLOYMENT_NAME,
        workspace="default"
    )
    
    elapsed = time.time() - start_time
    elapsed_str = f"{int(elapsed//60)}m {int(elapsed%60)}s"
    
    clear_output(wait=True)
    print(f"Baseline deployment: {status.status} | {elapsed_str}")
    
    if status.status == "READY":
        print("Baseline model ready!")
        if not client.models.wait_for_gateway(BASELINE_DEPLOYMENT_NAME, workspace="default", timeout=60):
            raise RuntimeError("Inference gateway did not become ready")
        break
    if status.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
        raise RuntimeError(f"Baseline deployment failed: {status.status}")
    if elapsed > TIMEOUT_MINUTES * 60:
        raise TimeoutError("Baseline deployment timeout")
    
    time.sleep(10)

# Run baseline ranking with the base model
BASE_MODEL_ID = "nvidia/llama-nemotron-embed-1b-v2"
# Get query embedding
query_response = client.inference.gateway.provider.post(
    "v1/embeddings",
    name=BASELINE_DEPLOYMENT_NAME,
    workspace="default",
    body={
        "model": BASE_MODEL_ID,
        "input": [DEMO_QUERY],
        "input_type": "query"
    }
)
base_query_emb = query_response["data"][0]["embedding"]
# Get document embeddings
doc_response = client.inference.gateway.provider.post(
    "v1/embeddings",
    name=BASELINE_DEPLOYMENT_NAME,
    workspace="default",
    body={
        "model": BASE_MODEL_ID,
        "input": DEMO_DOCS,
        "input_type": "passage"
    }
)
base_doc_embs = [d["embedding"] for d in doc_response["data"]]
# Calculate similarities and rank
scores = [(i, cosine_similarity(base_query_emb, base_doc_embs[i])) for i in range(len(DEMO_DOCS))]
BASELINE_RANKING = sorted(scores, key=lambda x: -x[1])
# Display baseline results
print(f"Query: \"{DEMO_QUERY}\"\n")
print("Base Model Ranking:")
print("-" * 55)
for rank, (idx, score) in enumerate(BASELINE_RANKING, 1):
    marker = " <-- relevant" if idx in DEMO_RELEVANT else ""
    print(f"  #{rank}  [{score:.3f}]  {DEMO_LABELS[idx]}{marker}")

# Delete baseline deployment to free GPU for training
print("Deleting baseline deployment to free GPU...")
client.inference.deployments.delete(name=BASELINE_DEPLOYMENT_NAME, workspace="default")
# Wait for deployment to be fully deleted before deleting config
print("Waiting for deployment deletion...")
while True:
    try:
        status = client.inference.deployments.retrieve(name=BASELINE_DEPLOYMENT_NAME, workspace="default")
        if status.status == "DELETED":
            break
        print(f"  Status: {status.status}")
        time.sleep(5)
    except Exception:
        # Deployment no longer exists
        break
# Now safe to delete the config
client.inference.deployment_configs.delete(name=BASELINE_DEPLOYMENT_CONFIG, workspace="default")
print("GPU freed. Proceed to fine-tune and improve these rankings.")

from pathlib import Path
from datasets import load_dataset
import json
# Configuration
DATASET_SIZE = 3000        # Number of triplets (increase for better results, max ~684K)
VALIDATION_SPLIT = 0.05   # 5% held out for validation
SEED = 42
DATASET_PATH = Path("embedding-dataset").absolute()
# Create directory
os.makedirs(DATASET_PATH, exist_ok=True)
# Download SPECTER dataset
print("Downloading SPECTER dataset...")
data = load_dataset("embedding-data/SPECTER")["train"].shuffle(seed=SEED).select(range(DATASET_SIZE))
# Split into train/validation
print("Splitting into train/validation...")
splits = data.train_test_split(test_size=VALIDATION_SPLIT, seed=SEED)
train_data = splits["train"]
validation_data = splits["test"]
# Convert to triplet JSONL format
print("Saving to JSONL...")
for name, dataset in [("training", train_data), ("validation", validation_data)]:
    with open(f"{DATASET_PATH}/{name}.jsonl", "w") as f:
        for row in dataset:
            # SPECTER format: row['set'] = [query, positive, negative]
            triplet = {
                "query": row["set"][0],
                "pos_doc": row["set"][1],
                "neg_doc": [row["set"][2]]  # List of negative documents
            }
            f.write(json.dumps(triplet) + "\n")
print(f"\nPrepared {len(train_data):,} training, {len(validation_data):,} validation samples")
print(f"\nExample triplet:")
print(f"  Query:    {train_data[0]['set'][0][:100]}...")
print(f"  Positive: {train_data[0]['set'][1][:100]}...")
print(f"  Negative: {train_data[0]['set'][2][:100]}...")

# Create fileset to store embedding training data
DATASET_NAME = "embedding-dataset"
try:
    client.files.filesets.create(
        workspace="default",
        name=DATASET_NAME,
        description="SPECTER embedding training data (scientific paper triplets)"
    )
    print(f"Created fileset: {DATASET_NAME}")
except ConflictError:
    print(f"Fileset '{DATASET_NAME}' already exists, continuing...")
# Upload training data files
client.files.upload(
    local_path=f"{DATASET_PATH}/",
    remote_path="",
    fileset=DATASET_NAME,
    workspace="default"
)
# Validate upload
print("\nUploaded files:")
print(json.dumps([f.model_dump() for f in client.files.list(fileset=DATASET_NAME, workspace="default").data], indent=2))

# Create secrets for model access
# Note: NGC_API_KEY secret was already created in the baseline step (Step 2)
HF_TOKEN = os.getenv("HF_TOKEN")
def create_or_get_secret(name: str, value: str | None, label: str):
    if not value:
        raise ValueError(f"{label} is not set")
    try:
        secret = client.secrets.create(
            name=name,
            workspace="default",
            value=value,
        )
        print(f"Created secret: {name}")
        return secret
    except ConflictError:
        print(f"Secret '{name}' already exists, continuing...")
        return client.secrets.retrieve(name=name, workspace="default")
# Public Hugging Face models need no token. Create a secret only when HF_TOKEN is set.
hf_secret = create_or_get_secret("hf-token", HF_TOKEN, "HF_TOKEN") if HF_TOKEN else None
if hf_secret:
    print(f"HF_TOKEN secret: {hf_secret.name}")
# NGC secret was already created in baseline step (Step 2), or use the platform default
if "NGC_SECRET_NAME" not in globals():
    NGC_SECRET_NAME = "ngc-api-key"
print(f"NGC_API_KEY secret: {NGC_SECRET_NAME}")

import time
from nemo_platform.types.files import HuggingfaceStorageConfigParam
HF_REPO_ID = "nvidia/llama-nemotron-embed-1b-v2"
MODEL_NAME = "nv-nemotron-embed-1b-base"
storage_kwargs = {
    "type": "huggingface",
    "repo_id": HF_REPO_ID,
    "repo_type": "model",
}
if hf_secret:
    storage_kwargs["token_secret"] = hf_secret.name
storage = HuggingfaceStorageConfigParam(**storage_kwargs)
try:
    base_model_fs = client.files.filesets.create(
        workspace="default",
        name=MODEL_NAME,
        description="NVIDIA Llama Nemotron Embed 1B v2 embedding model",
        storage=storage,
    )
except ConflictError as e:
    print(f"Base model fileset already exists. Skipping creation.")
    base_model_fs = client.files.filesets.retrieve(
        workspace="default",
        name=MODEL_NAME,
    )
# Create Model Entity referencing the FileSet
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:
    print(f"Base model already exists. Updating fileset if different.")
    base_model = client.models.update(
        workspace="default",
        name=MODEL_NAME,
        fileset=f"default/{MODEL_NAME}",
        trust_remote_code=True,
    )
print(f"\nBase model fileset: fileset://default/{base_model.name}")
print("\nBase model files:")
print(json.dumps([f.model_dump() for f in client.files.list(fileset=MODEL_NAME, workspace="default").data], indent=2))
# Wait for ModelSpec to be populated from the checkpoint
print("\nWaiting for ModelSpec to be populated...")
SPEC_TIMEOUT_SECONDS = 120
spec_start = time.time()
while not base_model.spec:
    if time.time() - spec_start > SPEC_TIMEOUT_SECONDS:
        raise TimeoutError(f"ModelSpec not populated within {SPEC_TIMEOUT_SECONDS} seconds")
    time.sleep(2)
    base_model = client.models.retrieve(
        workspace="default",
        name=MODEL_NAME,
    )
print(f"ModelSpec populated: {base_model.spec}")

training={
    "training_type": "sft",
    "finetuning_type": "lora_merged",
    "lora": {"rank": 16, "alpha": 32},
    "max_seq_length": 512,
}

import uuid
from nemo_automodel_plugin.schema import AutomodelJobInput
job_suffix = uuid.uuid4().hex[:4]
JOB_NAME = f"embedding-finetune-job-{job_suffix}"
OUTPUT_NAME = f"nv-embed-finetuned-{job_suffix}"
EPOCHS = 1
BATCH_SIZE = 128
LEARNING_RATE = 5e-6
MAX_SEQ_LENGTH = 512
spec = AutomodelJobInput(
    model=f"default/{base_model.name}",
    dataset={"training": f"default/{DATASET_NAME}"},
    training={
        "training_type": "sft",
        "finetuning_type": "all_weights",
        "max_seq_length": MAX_SEQ_LENGTH,
    },
    schedule={"epochs": EPOCHS},
    batch={"global_batch_size": BATCH_SIZE, "micro_batch_size": 1},
    optimizer={"learning_rate": LEARNING_RATE},
    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.job.name}")
print(f"Output model: {OUTPUT_NAME}")

import time
from IPython.display import clear_output
# Poll job status every 10 seconds until completed
while True:
    status = client.jobs.get_status(
        name=job.job.name,
        workspace="default"
    )
    
    clear_output(wait=True)
    print(f"Job Status: {status.model_dump_json(indent=2)}")
    # Extract training progress from nested steps structure
    step: int | None = None
    max_steps: int | None = None
    training_phase: str | None = None
    for job_step in status.steps or []:
        if job_step.name == "training":
            for task in job_step.tasks or []:
                task_details = task.status_details or {}
                step = task_details.get("step")
                max_steps = task_details.get("max_steps")
                training_phase = task_details.get("phase")
                break
            break
    if step is not None and max_steps is not None:
        progress_pct = (step / max_steps) * 100
        print(f"Training Progress: Step {step}/{max_steps} ({progress_pct:.1f}%)")
        if training_phase:
            print(f"Training Phase: {training_phase}")
    else:
        print("Training step not started yet or progress info not available")
    
    # Exit loop when job is completed (or failed/cancelled)
    if status.status in ("completed", "failed", "cancelled", "error"):
        print(f"\nJob finished with status: {status.status}")
        break
    
    time.sleep(10)
if status.status != "completed":
    raise RuntimeError(f"Training job finished with status: {status.status}")

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

# Create deployment config for embedding model
deploy_suffix = uuid.uuid4().hex[:4]
DEPLOYMENT_CONFIG_NAME = f"embedding-model-deployment-cfg-{deploy_suffix}"
DEPLOYMENT_NAME = f"embedding-model-deployment-{deploy_suffix}"
# Embedding NIM image
NIM_IMAGE = "nvcr.io/nim/nvidia/llama-nemotron-embed-1b-v2"
NIM_TAG = "1.13.0"  # Update if using newer NIM release
deployment_config = client.inference.deployment_configs.create(
    workspace="default",
    name=DEPLOYMENT_CONFIG_NAME,
    engine="nim",
    model_spec={
        "model_namespace": "default",
        "model_name": OUTPUT_NAME,
    },
    executor_config={
        "gpu": 1,
        "image_name": NIM_IMAGE,
        "image_tag": NIM_TAG,
    },
)
# Deploy model
deployment = client.inference.deployments.create(
    workspace="default",
    name=DEPLOYMENT_NAME,
    config=deployment_config.name
)
print(f"Deployment name: {deployment.name}")
print(f"Deployment status: {client.inference.deployments.retrieve(name=deployment.name, workspace='default').status}")

import time
from IPython.display import clear_output
# Poll deployment status every 15 seconds until ready
TIMEOUT_MINUTES = 30
start_time = time.time()
timeout_seconds = TIMEOUT_MINUTES * 60
print(f"Monitoring deployment '{deployment.name}'...")
print(f"Timeout: {TIMEOUT_MINUTES} minutes\n")
while True:
    deployment_status = client.inference.deployments.retrieve(
        name=deployment.name,
        workspace="default"
    )
    
    elapsed = time.time() - start_time
    elapsed_min = int(elapsed // 60)
    elapsed_sec = int(elapsed % 60)
    
    clear_output(wait=True)
    print(f"Deployment: {deployment.name}")
    print(f"Status: {deployment_status.status}")
    print(f"Elapsed time: {elapsed_min}m {elapsed_sec}s")
    
    # Check if deployment is ready
    if deployment_status.status == "READY":
        print("\nDeployment is ready!")
        if not client.models.wait_for_gateway(deployment.name, workspace="default", timeout=60):
            raise RuntimeError("Inference gateway did not become ready")
        break
    
    # Check for failure states
    if deployment_status.status in ("FAILED", "ERROR", "TERMINATED", "LOST"):
        raise RuntimeError(f"Deployment failed with status: {deployment_status.status}")
    
    # Check timeout
    if elapsed > timeout_seconds:
        raise TimeoutError(f"Deployment timeout after {TIMEOUT_MINUTES} minutes")
    
    time.sleep(15)

# Compare: same query, base model vs fine-tuned
# Using the same DEMO_QUERY and DEMO_DOCS from the baseline test
MODEL_ID = f"default/{OUTPUT_NAME}"
# Get query embedding from fine-tuned model
query_response = client.inference.gateway.provider.post(
    "v1/embeddings",
    name=deployment.name,
    workspace="default",
    body={
        "model": MODEL_ID,
        "input": [DEMO_QUERY],
        "input_type": "query"
    }
)
query_embedding = query_response["data"][0]["embedding"]
# Get document embeddings from fine-tuned model
doc_response = client.inference.gateway.provider.post(
    "v1/embeddings",
    name=deployment.name,
    workspace="default",
    body={
        "model": MODEL_ID,
        "input": DEMO_DOCS,
        "input_type": "passage"
    }
)
doc_embeddings = [d["embedding"] for d in doc_response["data"]]
# Calculate similarities and rank
scores = [(i, cosine_similarity(query_embedding, doc_embeddings[i])) for i in range(len(DEMO_DOCS))]
FINETUNED_RANKING = sorted(scores, key=lambda x: -x[1])
# Display side-by-side comparison
print(f"Query: \"{DEMO_QUERY}\"\n")
print(f"{'Rank':<6} {'Base Model':<30} {'Fine-tuned Model':<30}")
print("-" * 66)
for rank in range(len(DEMO_DOCS)):
    b_idx, b_score = BASELINE_RANKING[rank]
    f_idx, f_score = FINETUNED_RANKING[rank]
    
    b_label = f"{DEMO_LABELS[b_idx]} [{b_score:.3f}]" + (" *" if b_idx in DEMO_RELEVANT else "")
    f_label = f"{DEMO_LABELS[f_idx]} [{f_score:.3f}]" + (" *" if f_idx in DEMO_RELEVANT else "")
    
    print(f"#{rank+1:<5} {b_label:<30} {f_label:<30}")
print("\n* = relevant paper")
print("\nThe fine-tuned model pushes 'Random Forest' down and ranks CRF papers higher.")