com.sky.ondeviceagent / Runtime /AgentCore /Rag /E5TextEmbedder.cs
Sky-Kim's picture
Initial commit
2e7837a
Raw
History Blame Contribute Delete
2.36 kB
using System;
using System.Threading.Tasks;
using LightRAG.Core.Abstractions;
using Unity.InferenceEngine;
using SentisModels;
using OnDeviceAgent.Inference;
namespace OnDeviceAgent.AgentCore
{
// Framework adapter over SentisModels.E5Embedder: adds Android-safe loading, ITextEmbedder, and the LightRAG EmbeddingFunc bridge.
public sealed class E5TextEmbedder : ITextEmbedder
{
public const int FeatureDim = E5Embedder.FeatureDim;
const int MaxTokens = 512; // e5 context limit (used for the EmbeddingFunc token budget)
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));
// Sync EmbedText, not async: async readback can hang the serialized queue in batchmode (no SynchronizationContext).
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();
}
}