| using System; |
| using System.Threading.Tasks; |
| using LightRAG.Core.Abstractions; |
| using Unity.InferenceEngine; |
| using SentisModels; |
| using OnDeviceAgent.Inference; |
|
|
| namespace OnDeviceAgent.AgentCore |
| { |
| |
| public sealed class E5TextEmbedder : ITextEmbedder |
| { |
| public const int FeatureDim = E5Embedder.FeatureDim; |
| const int MaxTokens = 512; |
|
|
| readonly E5Embedder m_Core; |
|
|
| public bool Ready => m_Core.Ready; |
|
|
| public E5TextEmbedder(BackendType backend = BackendType.CPU, Action<string> log = null) |
| { |
| m_Core = new E5Embedder(backend, log); |
| } |
|
|
| public bool Load(string modelPath, string tokenizerPath) |
| { |
| var root = (System.IO.Path.GetDirectoryName(modelPath) ?? string.Empty).Replace('\\', '/'); |
| return m_Core.Load( |
| ModelRootProvisioner.EnsureModelRoot(root), |
| System.IO.Path.GetFileName(modelPath), |
| System.IO.Path.GetFileName(tokenizerPath)); |
| } |
|
|
| public float[] EmbedText(string text, bool isQuery = false) => m_Core.EmbedText(text, isQuery); |
|
|
| public Task<float[]> EmbedTextAsync(string text, bool isQuery = false) => m_Core.EmbedTextAsync(text, isQuery); |
|
|
| public EmbeddingFunc AsEmbeddingFunc(IUnityMainThreadDispatcher dispatcher) |
| { |
| if (dispatcher == null) throw new ArgumentNullException(nameof(dispatcher)); |
| |
| return new EmbeddingFunc( |
| FeatureDim, |
| (texts, context, ct) => dispatcher.RunOnMainAsync<float[][]>(() => |
| { |
| var isQuery = context == "query"; |
| var vectors = new float[texts.Count][]; |
| for (var i = 0; i < texts.Count; i++) |
| { |
| ct.ThrowIfCancellationRequested(); |
| vectors[i] = EmbedText(texts[i], isQuery); |
| } |
| return Task.FromResult(vectors); |
| }), |
| MaxTokenSize: MaxTokens, |
| ModelName: "multilingual-e5-small", |
| SupportsAsymmetric: true); |
| } |
|
|
| public void Dispose() => m_Core.Dispose(); |
| } |
| } |
|
|