ner

Documentation

Overview

Package ner provides ML-based named entity recognition for PII redaction. It uses ONNX Runtime (via purego) and a multilingual DistilBERT NER model to detect person names in freeform text.

The model bundle is downloaded separately via `hs pii-model install` and cached in the user's OS-specific cache directory. When the bundle is not present, the NER detector cannot be created and freeform text is hidden instead.

Index

• Constants
• func CacheDir() (string, error)
• func DecodeLogits(logits [][]float32, labels []string, seqLen int) []tokenTag
• func IsModelReady() bool
• func LoadLabels(configPath string) ([]string, error)
• func MergePersonSpans(tags []tokenTag, offsets [][2]int, text string) []pii.NameSpan
• func RemoveModel() error
• type Detector
  • func NewDetector() (*Detector, error)
  • func (d *Detector) Close()
  • func (d *Detector) DetectNames(text string) ([]pii.NameSpan, error)
• type Encoding
• type Paths
  • func EnsureModel(progress ProgressFunc) (*Paths, error)
  • func ModelPaths() (*Paths, error)
• type ProgressFunc
• type Runtime
  • func NewRuntime(paths *Paths) (*Runtime, error)
  • func (r *Runtime) Close()
  • func (r *Runtime) Run(inputIDs, attentionMask []int64) ([][]float32, error)
• type Tokenizer
  • func NewTokenizer(path string) (*Tokenizer, error)
  • func (t *Tokenizer) Encode(text string) (*Encoding, error)

Constants

const MinConfidence float32 = 0.7

MinConfidence is the default threshold for accepting a name span.

const ModelVersion = "0.2.0"

ModelVersion is the version tag used for bundle download URLs.

Functions

func CacheDir

func CacheDir() (string, error)

CacheDir returns the OS-specific cache directory for the PII model.

func DecodeLogits

func DecodeLogits(logits [][]float32, labels []string, seqLen int) []tokenTag

DecodeLogits converts raw logits to BIO tags via argmax.

func IsModelReady

func IsModelReady() bool

IsModelReady checks whether the model bundle is present and matches the expected version without downloading anything.

func LoadLabels

func LoadLabels(configPath string) ([]string, error)

LoadLabels reads the id2label mapping from config.json and returns an ordered slice where index = label ID.

func MergePersonSpans

func MergePersonSpans(tags []tokenTag, offsets [][2]int, text string) []pii.NameSpan

MergePersonSpans groups B-PER + I-PER sequences into NameSpans using offsets. Special tokens (offset 0,0) and non-PER labels are skipped.

func RemoveModel

func RemoveModel() error

RemoveModel deletes the cached model bundle.

Types

type Detector

type Detector struct {
	// contains filtered or unexported fields
}

Detector performs named entity recognition on text.

func NewDetector

func NewDetector() (*Detector, error)

NewDetector loads the ONNX Runtime and model from the cache directory. Returns an error if the model is not installed.

func (*Detector) Close

func (d *Detector) Close()

Close releases ONNX Runtime resources.

func (*Detector) DetectNames

func (d *Detector) DetectNames(text string) ([]pii.NameSpan, error)

DetectNames returns person name spans found in text. Long text is automatically chunked to stay within the model's 512-token limit.

type Encoding

type Encoding struct {
	IDs           []int64
	AttentionMask []int64
	Offsets       [][2]int // (start, end) byte offsets into original text per token
	Tokens        []string
}

Encoding holds tokenizer output.

type Paths

type Paths struct {
	RuntimeLib    string // platform-specific ONNX Runtime shared lib
	ModelONNX     string // model_quantized.onnx
	TokenizerJSON string // tokenizer.json
	ConfigJSON    string // config.json
}

Paths holds resolved file paths for the model bundle.

func EnsureModel

func EnsureModel(progress ProgressFunc) (*Paths, error)

EnsureModel downloads and extracts the model bundle if not present.

func ModelPaths

func ModelPaths() (*Paths, error)

ModelPaths returns resolved paths if the model is installed, or an error.

type ProgressFunc

type ProgressFunc func(bytesRead, totalBytes int64)

ProgressFunc reports download progress (bytesRead, totalBytes). totalBytes may be -1 if unknown.

type Runtime

type Runtime struct {
	// contains filtered or unexported fields
}

Runtime wraps an ONNX Runtime session for NER inference.

func NewRuntime

func NewRuntime(paths *Paths) (*Runtime, error)

NewRuntime loads the ONNX Runtime shared library and model.

func (*Runtime) Close

func (r *Runtime) Close()

Close releases ONNX Runtime resources.

func (*Runtime) Run

func (r *Runtime) Run(inputIDs, attentionMask []int64) ([][]float32, error)

Run executes inference on the given input_ids and attention_mask. Returns logits as [seqLen][numLabels].

type Tokenizer

type Tokenizer struct {
	// contains filtered or unexported fields
}

Tokenizer wraps a HuggingFace tokenizer loaded from tokenizer.json.

func NewTokenizer

func NewTokenizer(path string) (*Tokenizer, error)

NewTokenizer loads a tokenizer from a tokenizer.json file.

func (*Tokenizer) Encode

func (t *Tokenizer) Encode(text string) (*Encoding, error)

Encode tokenizes text and returns IDs, attention mask, and offsets.