ossamma-ner-onnx

OssammaNER - Named Entity Recognition Model

OssammaNER is a 75M parameter NER model using the Ossamma architecture (Oscillatory State Space Attention Masked Mixer Architecture).

Model Details

• Parameters: ~75M
• Architecture: Custom Ossamma blocks with:
  • DLinOSS (Damped Linear Oscillator State Space Model)
  • Linear Attention with positional embeddings
  • Sliding Window Attention with SigSoftmax
  • SwiGLU FFN
  • CRF decoder for sequence labeling
• Trained on: MultiNERD + Few-NERD (300K+ examples)
• Labels: 17 (8 entity types with BIO tagging)

Entity Types

Usage

ONNX Inference (Python)

import onnxruntime as ort
import numpy as np

# Load model
session = ort.InferenceSession("model.onnx")

# Tokenize input (requires compatible tokenizer)
token_ids = np.array([[101, 7592, 1010, 2088, 999, 102]])  # Example

# Run inference
emissions, boundary_logits = session.run(None, {"token_ids": token_ids})

# Get predictions
predictions = np.argmax(emissions, axis=-1)

Julia Inference (Native)

For full model fidelity including CRF decoding, use the Julia implementation:

using Ossamma
using Ossamma.NER: OssammaNER, predict_labels, extract_entities

# Load model
model = OssammaNER("configs/ner_production_110m.toml")
checkpoint = load_checkpoint("checkpoint_best.jls")
ps, st = checkpoint.params, checkpoint.state

# Inference
tokens = [101, 7592, 1010, 2088, 999, 102]
labels = predict_labels(model, ps, st, tokens)
entities = extract_entities(tokens, labels)

Training Details

• Training Steps: 50,000
• Batch Size: 8 (effective 64 with gradient accumulation)
• Learning Rate: 2e-4 with cosine decay
• Hardware: NVIDIA RTX 5090 (32GB VRAM)
• Training Time: ~5 hours

Model Architecture

OssammaNER(
  TokenEmbedding: 32000 × 640
  PositionEmbedding: 256 × 640
  TimeEmbedding: Sinusoidal (192 dim)

  Blocks × 10:
    TimeConditionedLayerNorm
    GLU Branch:
      - Dense(640 → 1280)
      - LinearAttention (10 heads)
      - DLinOSS (oscillator SSM)
      - Dense(640 → 640)
    Local Branch:
      - SWAttention (window=32)
      - Input Gate
    Adaptive Mixing (α-blending)
    SwiGLU FFN (640 → 960 → 640)
    OutputLayerNorm

  ClassificationHead: LayerNorm → Dropout → Dense(640 → 17)
  BoundaryHead: LayerNorm → Dense(640 → 2)
  CRF: LinearChainCRF (17 labels)
)

Limitations

• This ONNX model uses a simplified transformer architecture for compatibility
• For best performance, use the native Julia implementation with the .jls checkpoint
• The CRF decoder is not included in the ONNX model (use Viterbi decoding separately)

Citation

@software{ossamma_ner,
  title={OssammaNER: Oscillatory State Space Attention for Named Entity Recognition},
  author={Biotz},
  year={2024},
  url={https://huggingface.co/Biotz/ossamma-ner}
}

License

MIT License