llm-pii-poc/src/PiiRedaction.Infrastructure/Onnx/OnnxTokenClassifierRunner.cs

using Microsoft.Extensions.Logging;
using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using PiiRedaction.Core.Models;

namespace PiiRedaction.Infrastructure.Onnx;

/// <summary>
/// Shared ONNX token-classification inference and BIO decoding for NER models.
/// </summary>
public sealed class OnnxTokenClassifierRunner : IDisposable
{
    private const int MaxSequenceLength = 128;

    private readonly ITokenClassifierEncoder _encoder;
    private readonly NerLabelConfig _labelConfig;
    private readonly string[] _labels;
    private readonly ILogger _logger;
    private readonly string _modelPath;
    private InferenceSession? _session;

    public OnnxTokenClassifierRunner(
        string modelPath,
        ITokenClassifierEncoder encoder,
        NerLabelConfig labelConfig,
        string[] labels,
        ILogger logger)
    {
        _modelPath = modelPath;
        _encoder = encoder;
        _labelConfig = labelConfig;
        _labels = labels;
        _logger = logger;
        _session = TryCreateSession();
    }

    public bool IsAvailable => _session is not null && _encoder.IsAvailable && _labels.Length > 0;

    public IReadOnlyList<PiiEntity> PredictEntities(string text)
    {
        if (!IsAvailable || _session is null)
        {
            return [];
        }

        var encoded = _encoder.Encode(text, MaxSequenceLength);
        if (encoded is null)
        {
            return [];
        }

        var predictedLabelIds = RunInference(encoded);
        return DecodePersonEntities(text, predictedLabelIds, encoded);
    }

    private int[] RunInference(EncodedSequence encoded)
    {
        var inputIdsTensor = CreateTensor(encoded.InputIds, encoded.SequenceLength);
        var attentionMaskTensor = CreateTensor(encoded.AttentionMask, encoded.SequenceLength);
        var inputs = new List<NamedOnnxValue>
        {
            NamedOnnxValue.CreateFromTensor(
                _session!.InputMetadata.Keys.First(key => key.Contains("input_ids", StringComparison.OrdinalIgnoreCase)),
                inputIdsTensor),
            NamedOnnxValue.CreateFromTensor(
                _session.InputMetadata.Keys.First(key => key.Contains("attention_mask", StringComparison.OrdinalIgnoreCase)),
                attentionMaskTensor)
        };

        var tokenTypeInputName = _session.InputMetadata.Keys.FirstOrDefault(key =>
            key.Contains("token_type", StringComparison.OrdinalIgnoreCase));
        if (tokenTypeInputName is not null)
        {
            inputs.Add(NamedOnnxValue.CreateFromTensor(
                tokenTypeInputName,
                CreateTensor(encoded.TokenTypeIds, encoded.SequenceLength)));
        }

        using var results = _session.Run(inputs);
        var outputName = _session.OutputMetadata.Keys.FirstOrDefault(key =>
            key.Contains("logits", StringComparison.OrdinalIgnoreCase))
            ?? results.First().Name;
        var logits = results.First(result => result.Name == outputName).AsTensor<float>();
        var numLabels = _labels.Length;
        var predictions = new int[encoded.SequenceLength];

        for (var tokenIndex = 0; tokenIndex < encoded.SequenceLength; tokenIndex++)
        {
            var bestLabel = 0;
            var bestScore = float.MinValue;

            for (var labelIndex = 0; labelIndex < numLabels; labelIndex++)
            {
                var score = logits[0, tokenIndex, labelIndex];
                if (score > bestScore)
                {
                    bestScore = score;
                    bestLabel = labelIndex;
                }
            }

            predictions[tokenIndex] = bestLabel;
        }

        return predictions;
    }

    private static DenseTensor<long> CreateTensor(long[] values, int sequenceLength)
    {
        var tensor = new DenseTensor<long>([1, sequenceLength]);
        for (var i = 0; i < sequenceLength; i++)
        {
            tensor[0, i] = values[i];
        }

        return tensor;
    }

    private IReadOnlyList<PiiEntity> DecodePersonEntities(
        string text,
        int[] predictedLabelIds,
        EncodedSequence encoded)
    {
        var entities = new List<PiiEntity>();
        int? entityStart = null;
        int? entityEnd = null;

        void FlushEntity()
        {
            if (!entityStart.HasValue || !entityEnd.HasValue || entityEnd.Value <= entityStart.Value)
            {
                entityStart = null;
                entityEnd = null;
                return;
            }

            var value = text[entityStart.Value..entityEnd.Value];
            if (!string.IsNullOrWhiteSpace(value))
            {
                entities.Add(new PiiEntity(
                    PiiEntityType.Person,
                    value,
                    entityStart.Value,
                    entityEnd.Value - entityStart.Value,
                    PiiDetectionSource.Ner));
            }

            entityStart = null;
            entityEnd = null;
        }

        for (var i = 0; i < encoded.SequenceLength; i++)
        {
            if (_encoder.IsSpecialToken(encoded.TokenIds[i]))
            {
                FlushEntity();
                continue;
            }

            var label = _labels[predictedLabelIds[i]];
            var (start, end) = encoded.Offsets[i];
            var hasOffset = end > start;

            if (!_labelConfig.IsPersonLabel(label))
            {
                FlushEntity();
                continue;
            }

            if (_labelConfig.IsBeginLabel(label))
            {
                FlushEntity();
                if (hasOffset)
                {
                    entityStart = start;
                    entityEnd = end;
                }
            }
            else if (_labelConfig.IsInsideLabel(label))
            {
                if (!entityStart.HasValue && hasOffset)
                {
                    entityStart = start;
                    entityEnd = end;
                }
                else if (hasOffset)
                {
                    entityEnd = Math.Max(entityEnd ?? end, end);
                }
            }
        }

        FlushEntity();
        return entities;
    }

    private InferenceSession? TryCreateSession()
    {
        if (!File.Exists(_modelPath))
        {
            _logger.LogWarning(
                "ONNX NER model not found at {ModelPath}. Person-name detection will return no results.",
                _modelPath);
            return null;
        }

        if (!_encoder.IsAvailable || _labels.Length == 0)
        {
            _logger.LogWarning(
                "Tokenizer or label map missing for ONNX NER model at {ModelPath}.",
                _modelPath);
            return null;
        }

        try
        {
            var session = new InferenceSession(_modelPath);
            _logger.LogInformation("ONNX NER model loaded from {ModelPath}.", _modelPath);
            return session;
        }
        catch (Exception ex)
        {
            _logger.LogError(ex, "Failed to load ONNX NER model from {ModelPath}.", _modelPath);
            return null;
        }
    }

    public void Dispose() => _session?.Dispose();
}