pplx-embed-0.6b-nemotron

Usage with mlx-raclate

This model can be used with mlx-raclate for native inference on Apple Silicon.

from mlx_raclate.utils.utils import load
from mlx_raclate.utils.token_classification import (
    postprocess_token_classification_output,
    viterbi_transition_biases_from_calibration,
)

# Load model and tokenizer
model_path = "PITTI/pplx-embed-0.6b-nemotron"  
model, tokenizer = load(
    model_path,
    pipeline="token-classification"
)

# Prepare input texts
texts = ['John works at Apple in California.', 'Microsoft was founded by Bill Gates.']

# Tokenize
max_length = getattr(model.config, "max_position_embeddings", 512)
tokens = tokenizer._tokenizer(
    texts,
    return_tensors="mlx",
    padding=True,
    truncation=True,
    max_length=max_length,
    return_offsets_mapping=True,
)
offset_mapping = tokens.pop("offset_mapping")

# Run inference
outputs = model(
    input_ids=tokens["input_ids"],
    attention_mask=tokens["attention_mask"],
    return_dict=True
)

# Get predictions
logits = outputs["logits"]
id2label = model.config.id2label
transition_biases = viterbi_transition_biases_from_calibration(
    getattr(model, "viterbi_calibration", None)
)
processed = postprocess_token_classification_output(
    logits=logits,
    probabilities=outputs["probabilities"],
    id2label=id2label,
    texts=texts,
    offsets=offset_mapping.tolist(),
    transition_biases=transition_biases,
)

# Process and print grouped spans
for i, text in enumerate(texts):
    print(f"Text: {text}")
    print("Grouped spans:")
    for span in processed["grouped_spans"][i]:
        print(f"  {span['entity_group']}: {span['word']!r} [{span['start']}, {span['end']}] score={span['score']:.3f}")
    print()

Model Details

• Base Model: perplexity-ai/pplx-embed-v1-0.6b
• Pipeline: token-classification
• Framework: mlx-raclate (MLX) or transformers

Inspiration

OpenMed/privacy-filter-nemotron, an amazing project led by Maziyar Panahi