Biotz/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

Label	Description
ART	Artworks, creative works
BUILDING	Buildings, structures
EVENT	Events, occurrences
LOCATION	Places, geographic entities
ORGANIZATION	Companies, institutions
OTHER	Miscellaneous entities
PERSON	People, characters
PRODUCT	Products, items

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