llm-pii-poc/tests/TestSupport.Shared/RealRoutingNerModelFixture.cs

using Microsoft.Extensions.Logging.Abstractions;
using Microsoft.Extensions.Options;
using PiiRedaction.Core.Configuration;
using PiiRedaction.Infrastructure.Onnx;

namespace PiiRedaction.Tests.Shared;

/// <summary>
/// Loads English and Tamil ONNX runners and exposes a <see cref="RoutingOnnxNerModelRunner"/>.
/// Skips when either model is missing or cannot be loaded.
/// </summary>
public abstract class RealRoutingNerModelFixture
{
    protected RoutingOnnxNerModelRunner Runner { get; private set; } = null!;

    protected EnglishOnnxNerRunner EnglishRunner { get; private set; } = null!;

    protected TamilOnnxNerRunner TamilRunner { get; private set; } = null!;

    [OneTimeSetUp]
    public void OneTimeSetUpRoutingModels()
    {
        var englishPath = RealNerModelPaths.ResolveRepoModelPath();
        if (!File.Exists(englishPath))
        {
            Assert.Ignore(RealNerModelPaths.ModelMissingMessage);
        }

        var tamilPath = RealTamilNerModelPaths.ResolveRepoModelPath();
        if (!File.Exists(tamilPath))
        {
            Assert.Ignore(RealTamilNerModelPaths.ModelMissingMessage);
        }

        var options = Options.Create(new PiiRedactionOptions
        {
            EnglishOnnxModelPath = englishPath,
            OnnxModelPath = englishPath,
            TamilOnnxModelPath = tamilPath,
            EnableTamilNer = true
        });

        EnglishRunner = new EnglishOnnxNerRunner(options, NullLogger<EnglishOnnxNerRunner>.Instance);
        TamilRunner = new TamilOnnxNerRunner(options, NullLogger<TamilOnnxNerRunner>.Instance);

        if (!EnglishRunner.IsModelAvailable)
        {
            EnglishRunner.Dispose();
            TamilRunner.Dispose();
            Assert.Ignore(RealNerModelPaths.ModelMissingMessage);
        }

        if (!TamilRunner.IsModelAvailable)
        {
            EnglishRunner.Dispose();
            TamilRunner.Dispose();
            Assert.Ignore(RealTamilNerModelPaths.ModelMissingMessage);
        }

        Runner = new RoutingOnnxNerModelRunner(EnglishRunner, TamilRunner, options);
    }

    [OneTimeTearDown]
    public void OneTimeTearDownRoutingModels()
    {
        EnglishRunner?.Dispose();
        TamilRunner?.Dispose();
    }
}