feat: initial commit
This commit is contained in:
Menci
@@ -0,0 +1,72 @@
|
||||
using MaigoLabs.NeedLe.Common;
|
||||
using MaigoLabs.NeedLe.Common.Extensions;
|
||||
using MaigoLabs.NeedLe.Common.Types;
|
||||
using MaigoLabs.NeedLe.Searcher.Trie;
|
||||
|
||||
namespace MaigoLabs.NeedLe.Searcher;
|
||||
|
||||
public class LoadedInvertedIndex
|
||||
{
|
||||
public class TokenDocumentReference
|
||||
{
|
||||
public required int DocumentId { get; set; }
|
||||
public required OffsetSpan[] Offsets { get; set; }
|
||||
}
|
||||
|
||||
public class TokenDefinitionExtended : TokenDefinition
|
||||
{
|
||||
public required TokenDocumentReference[] References { get; set; }
|
||||
}
|
||||
|
||||
public class TypedTries
|
||||
{
|
||||
public required TrieNode Romaji { get; set; }
|
||||
public required TrieNode Kana { get; set; }
|
||||
public required TrieNode Other { get; set; }
|
||||
}
|
||||
|
||||
public required string[] Documents { get; set; }
|
||||
public required int[][] DocumentCodePoints { get; set; }
|
||||
public required TokenDefinitionExtended[] TokenDefinitions { get; set; }
|
||||
public required TypedTries Tries { get; set; }
|
||||
}
|
||||
|
||||
public class InvertedIndexLoader
|
||||
{
|
||||
public static LoadedInvertedIndex Load(CompressedInvertedIndex compressed)
|
||||
{
|
||||
var documents = compressed.documents;
|
||||
var documentCodePoints = documents.Select(document => document.ToCodePoints().ToArray()).ToArray();
|
||||
|
||||
var romajiTrie = TrieDeserializer.Deserialize(compressed.tries.romaji);
|
||||
var kanaTrie = TrieDeserializer.Deserialize(compressed.tries.kana);
|
||||
var otherTrie = TrieDeserializer.Deserialize(compressed.tries.other);
|
||||
|
||||
var tokenCodePoints = romajiTrie.TokenCodePoints.Concat(kanaTrie.TokenCodePoints).Concat(otherTrie.TokenCodePoints)
|
||||
.ToDictionary(entry => entry.Key, entry => entry.Value);
|
||||
var tokenDefinitions = compressed.tokenTypes.Select((type, index) => new LoadedInvertedIndex.TokenDefinitionExtended
|
||||
{
|
||||
Id = index, Type = (TokenType)type, Text = tokenCodePoints[index].ToUtf32String(),
|
||||
CodePointLength = tokenCodePoints[index].Length,
|
||||
References = compressed.tokenReferences[index].Select(data => new LoadedInvertedIndex.TokenDocumentReference
|
||||
{
|
||||
DocumentId = data[0],
|
||||
Offsets = Enumerable.Range(0, data.Length / 2)
|
||||
.Select(i => new OffsetSpan { Start = data[i * 2 + 1], End = data[i * 2 + 2] }).ToArray(),
|
||||
}).ToArray(),
|
||||
}).ToArray();
|
||||
|
||||
return new LoadedInvertedIndex
|
||||
{
|
||||
Documents = documents,
|
||||
DocumentCodePoints = documentCodePoints,
|
||||
TokenDefinitions = tokenDefinitions,
|
||||
Tries = new LoadedInvertedIndex.TypedTries
|
||||
{
|
||||
Romaji = romajiTrie.Root,
|
||||
Kana = kanaTrie.Root,
|
||||
Other = otherTrie.Root,
|
||||
},
|
||||
};
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,270 @@
|
||||
using MaigoLabs.NeedLe.Common;
|
||||
using MaigoLabs.NeedLe.Common.Extensions;
|
||||
using MaigoLabs.NeedLe.Common.Types;
|
||||
|
||||
namespace MaigoLabs.NeedLe.Searcher;
|
||||
|
||||
public class SearchResultToken
|
||||
{
|
||||
public required TokenDefinition Definition { get; set; }
|
||||
public required OffsetSpan DocumentOffset { get; set; }
|
||||
public required OffsetSpan InputOffset { get; set; }
|
||||
public required bool IsTokenPrefixMatching { get; set; }
|
||||
}
|
||||
|
||||
public class SearchResult
|
||||
{
|
||||
public required int DocumentId { get; set; }
|
||||
public required string DocumentText { get; set; }
|
||||
public required int[] DocumentCodePoints { get; set; }
|
||||
public required SearchResultToken[] Tokens { get; set; }
|
||||
public required int PrefixMatchCount { get; set; }
|
||||
public required int RangeCount { get; set; }
|
||||
public required double MatchRatio { get; set; }
|
||||
public required int MatchRatioLevel { get; set; }
|
||||
}
|
||||
|
||||
public static class InvertedIndexSearcher
|
||||
{
|
||||
public abstract class ComparableStateBase<T> : IComparable<T>
|
||||
where T : ComparableStateBase<T>
|
||||
{
|
||||
protected abstract int GetRangeCount();
|
||||
protected abstract int GetPrefixMatchCount();
|
||||
protected abstract OffsetSpan GetFirstTokenDocumentOffset();
|
||||
protected abstract OffsetSpan GetLastTokenDocumentOffset();
|
||||
protected virtual SearchResultToken? GetLastToken() => null; // Not on intermediate results
|
||||
protected virtual int? GetMatchRatioLevel() => null; // Not on intermediate/candidate results
|
||||
protected abstract double GetMatchRatio();
|
||||
protected virtual int FallbackCompareTo(T other) => 0; // Called when all other comparisons are equal
|
||||
|
||||
public int CompareTo(T other)
|
||||
{
|
||||
// Prefer matches that not relying on end-of-input loose matching (full match over prefix match)
|
||||
SearchResultToken? aLastToken = GetLastToken(), bLastToken = other.GetLastToken();
|
||||
if (aLastToken != null && bLastToken != null)
|
||||
{
|
||||
var aDidPrefixMatchByTokenType = aLastToken.IsTokenPrefixMatching && tokenTypePrefixMatchingPolicy[aLastToken.Definition.Type] == TokenTypePrefixMatchingPolicy.AllowOnlyAtInputEnd;
|
||||
var bDidPrefixMatchByTokenType = bLastToken.IsTokenPrefixMatching && tokenTypePrefixMatchingPolicy[bLastToken.Definition.Type] == TokenTypePrefixMatchingPolicy.AllowOnlyAtInputEnd;
|
||||
if (aDidPrefixMatchByTokenType != bDidPrefixMatchByTokenType) return aDidPrefixMatchByTokenType ? 1 : -1;
|
||||
}
|
||||
|
||||
// Prefer results that matched fewer discontinuous ranges over more
|
||||
int aRangeCount = GetRangeCount(), bRangeCount = other.GetRangeCount();
|
||||
if (aRangeCount != bRangeCount) return aRangeCount - bRangeCount;
|
||||
|
||||
// Prefer results that matches first token in document earlier over later
|
||||
OffsetSpan aFirstTokenDocumentOffset = GetFirstTokenDocumentOffset(), bFirstTokenDocumentOffset = other.GetFirstTokenDocumentOffset();
|
||||
if (aFirstTokenDocumentOffset.Start != bFirstTokenDocumentOffset.Start) return aFirstTokenDocumentOffset.Start - bFirstTokenDocumentOffset.Start;
|
||||
|
||||
// Prefer results that has higher match ratio (but don't distinguish similar ratios, so we introduced `matchRatioLevel`)
|
||||
int? aMatchRatioLevel = GetMatchRatioLevel(), bMatchRatioLevel = other.GetMatchRatioLevel();
|
||||
if (aMatchRatioLevel != null && bMatchRatioLevel != null)
|
||||
{
|
||||
if (aMatchRatioLevel.Value != bMatchRatioLevel.Value) return bMatchRatioLevel.Value - aMatchRatioLevel.Value;
|
||||
}
|
||||
|
||||
// Prefer results that last token occurred earlier (if same, ended earlier) in the document over later
|
||||
OffsetSpan aLastTokenDocumentOffset = GetLastTokenDocumentOffset(), bLastTokenDocumentOffset = other.GetLastTokenDocumentOffset();
|
||||
if (aLastTokenDocumentOffset.Start != bLastTokenDocumentOffset.Start) return aLastTokenDocumentOffset.Start - bLastTokenDocumentOffset.Start;
|
||||
if (aLastTokenDocumentOffset.End != bLastTokenDocumentOffset.End) return aLastTokenDocumentOffset.End - bLastTokenDocumentOffset.End;
|
||||
|
||||
// Prefer results that has higher match ratio (precisely)
|
||||
double aMatchRatio = GetMatchRatio(), bMatchRatio = other.GetMatchRatio();
|
||||
if (aMatchRatio != bMatchRatio) return bMatchRatio < aMatchRatio ? -1 : bMatchRatio > aMatchRatio ? 1 : 0;
|
||||
|
||||
return FallbackCompareTo(other);
|
||||
}
|
||||
}
|
||||
|
||||
public class IntermediateResult : ComparableStateBase<IntermediateResult>
|
||||
{
|
||||
public required IntermediateResult? PreviousState { get; init; }
|
||||
public required OffsetSpan FirstTokenDocumentOffset { get; init; }
|
||||
public required int RangeCount { get; init; }
|
||||
public required int TokenCount { get; init; }
|
||||
public required int PrefixMatchCount { get; init; }
|
||||
public required double MatchedTokenLength { get; init; }
|
||||
public required int TokenId { get; init; }
|
||||
public required OffsetSpan DocumentOffset { get; init; }
|
||||
public required OffsetSpan InputOffset { get; init; }
|
||||
public required bool IsTokenPrefixMatching { get; init; }
|
||||
|
||||
protected override int GetRangeCount() => RangeCount;
|
||||
protected override int GetPrefixMatchCount() => PrefixMatchCount;
|
||||
protected override OffsetSpan GetFirstTokenDocumentOffset() => FirstTokenDocumentOffset;
|
||||
protected override OffsetSpan GetLastTokenDocumentOffset() => DocumentOffset;
|
||||
protected override double GetMatchRatio() => MatchedTokenLength; // No need to divide document length since intermediate results are for same document
|
||||
}
|
||||
|
||||
public class CandidateResult : ComparableStateBase<CandidateResult>
|
||||
{
|
||||
public required SearchResultToken[] Tokens { get; init; }
|
||||
public required int PrefixMatchCount { get; init; }
|
||||
public required double MatchedTokenLength { get; init; }
|
||||
public required int RangeCount { get; init; }
|
||||
|
||||
protected override int GetRangeCount() => RangeCount;
|
||||
protected override int GetPrefixMatchCount() => PrefixMatchCount;
|
||||
protected override OffsetSpan GetFirstTokenDocumentOffset() => Tokens[0].DocumentOffset;
|
||||
protected override OffsetSpan GetLastTokenDocumentOffset() => Tokens[^1].DocumentOffset;
|
||||
protected override SearchResultToken? GetLastToken() => Tokens[^1];
|
||||
protected override double GetMatchRatio() => MatchedTokenLength; // No need to divide document length since intermediate results are for same document
|
||||
}
|
||||
|
||||
public class FinalResult : ComparableStateBase<FinalResult>
|
||||
{
|
||||
public required SearchResult Result { get; init; }
|
||||
|
||||
protected override int GetRangeCount() => Result.RangeCount;
|
||||
protected override int GetPrefixMatchCount() => Result.PrefixMatchCount;
|
||||
protected override OffsetSpan GetFirstTokenDocumentOffset() => Result.Tokens[0].DocumentOffset;
|
||||
protected override OffsetSpan GetLastTokenDocumentOffset() => Result.Tokens[^1].DocumentOffset;
|
||||
protected override SearchResultToken? GetLastToken() => Result.Tokens[^1];
|
||||
protected override double GetMatchRatio() => Result.MatchRatio;
|
||||
protected override int? GetMatchRatioLevel() => Result.MatchRatioLevel;
|
||||
protected override int FallbackCompareTo(FinalResult other) => string.Compare(Result.DocumentText, other.Result.DocumentText, StringComparison.InvariantCulture);
|
||||
}
|
||||
|
||||
private static bool IsIgnorableCodePoint(int codePoint) => CommonUtils.IsWhitespace(codePoint) || codePoint == 0x3099 || codePoint == 0x309A;
|
||||
|
||||
public enum TokenTypePrefixMatchingPolicy {
|
||||
AlwaysAllow,
|
||||
NeverAllow,
|
||||
AllowOnlyAtInputEnd,
|
||||
}
|
||||
|
||||
private static Dictionary<TokenType, TokenTypePrefixMatchingPolicy> tokenTypePrefixMatchingPolicy = new()
|
||||
{
|
||||
[TokenType.Romaji] = TokenTypePrefixMatchingPolicy.NeverAllow,
|
||||
[TokenType.Kana] = TokenTypePrefixMatchingPolicy.AlwaysAllow,
|
||||
// These token types are in an "other" Trie
|
||||
[TokenType.Han] = TokenTypePrefixMatchingPolicy.AllowOnlyAtInputEnd, // No effect because always 1 code point
|
||||
[TokenType.Pinyin] = TokenTypePrefixMatchingPolicy.AllowOnlyAtInputEnd,
|
||||
[TokenType.Raw] = TokenTypePrefixMatchingPolicy.AllowOnlyAtInputEnd, // No effect because always 1 code point
|
||||
};
|
||||
|
||||
private static bool ShouldAllowPrefixMatching(TokenType tokenType, bool isAtInputEnd) =>
|
||||
tokenTypePrefixMatchingPolicy[tokenType] == TokenTypePrefixMatchingPolicy.AlwaysAllow ||
|
||||
(tokenTypePrefixMatchingPolicy[tokenType] != TokenTypePrefixMatchingPolicy.NeverAllow && isAtInputEnd);
|
||||
|
||||
private static bool HasNonEmptyCharacters(int[] documentCodePoints, int start, int end) =>
|
||||
start != end && !documentCodePoints.Skip(start).Take(end - start).All(CommonUtils.IsWhitespace);
|
||||
|
||||
public static SearchResult[] Search(LoadedInvertedIndex invertedIndex, string text)
|
||||
{
|
||||
var documents = invertedIndex.Documents;
|
||||
var documentCodePoints = invertedIndex.DocumentCodePoints;
|
||||
var tokenDefinitions = invertedIndex.TokenDefinitions;
|
||||
var tries = invertedIndex.Tries;
|
||||
|
||||
var codePoints = text.ToCodePoints().Select(CommonNormalization.NormalizeCodePoint).Select(CommonNormalization.ToKatakana).ToArray();
|
||||
// dp[i] = docId => end => IntermediateResult, starts from dp[-1] (l === 0), ends at dp[N - 1] (r === N - 1)
|
||||
var dp = Enumerable.Range(0, codePoints.Length).Select(l => new Dictionary<int, Dictionary<int, IntermediateResult>>()).ToArray();
|
||||
for (var l = 0; l < codePoints.Length; l++)
|
||||
{
|
||||
if (l != 0 && dp[l - 1].Count == 0) continue; // No documents match input from beginning to this position
|
||||
var romajiNode = tries.Romaji;
|
||||
var kanaNode = tries.Kana;
|
||||
var otherNode = tries.Other;
|
||||
for (var r = l; r < codePoints.Length && (romajiNode != null || kanaNode != null || otherNode != null); r++) // [l, r]
|
||||
{
|
||||
var codePoint = codePoints[r];
|
||||
romajiNode = romajiNode.TraverseStep(codePoint, IsIgnorableCodePoint(codePoint));
|
||||
kanaNode = kanaNode.TraverseStep(codePoint, IsIgnorableCodePoint(codePoint));
|
||||
otherNode = otherNode.TraverseStep(codePoint, IsIgnorableCodePoint(codePoint));
|
||||
var reachingInputEnd = r == codePoints.Length - 1;
|
||||
HashSet<int> matchingTokenIds =
|
||||
[
|
||||
// Allow suffix matching of romaji/other tokens if we're at the end of the input
|
||||
.. romajiNode.GetTokenIds(ShouldAllowPrefixMatching(TokenType.Romaji, reachingInputEnd)),
|
||||
.. kanaNode.GetTokenIds(ShouldAllowPrefixMatching(TokenType.Kana, reachingInputEnd)),
|
||||
.. otherNode.GetTokenIds(reachingInputEnd),
|
||||
];
|
||||
foreach (var tokenId in matchingTokenIds) foreach (var reference in tokenDefinitions[tokenId].References)
|
||||
{
|
||||
var isTokenPrefixMatching = !romajiNode.IsTokenExactMatch(tokenId) && !kanaNode.IsTokenExactMatch(tokenId) && !otherNode.IsTokenExactMatch(tokenId);
|
||||
var previousMatchesOfDocument = l != 0 && dp[l - 1].TryGetValue(reference.DocumentId, out var previousMatches) ? previousMatches : null;
|
||||
if (l != 0 && previousMatchesOfDocument == null) continue;
|
||||
foreach (var documentOffset in reference.Offsets)
|
||||
{
|
||||
int currentStart = documentOffset.Start, currentEnd = documentOffset.End;
|
||||
if (l == 0) ContributeNextMatchingState(null);
|
||||
else foreach (var (previousEnd, previousMatch) in previousMatchesOfDocument!) if (currentStart >= previousEnd) ContributeNextMatchingState(previousMatch);
|
||||
void ContributeNextMatchingState(IntermediateResult? previousState)
|
||||
{
|
||||
var nextMatchingMap = dp[r];
|
||||
if (!nextMatchingMap.TryGetValue(reference.DocumentId, out var nextMatches)) nextMatches = nextMatchingMap[reference.DocumentId] = [];
|
||||
var oldResult = nextMatches.TryGetValue(currentEnd, out var result) ? result : null;
|
||||
var inputOffset = new OffsetSpan { Start = l, End = r + 1 };
|
||||
var newResult = new IntermediateResult
|
||||
{
|
||||
PreviousState = previousState,
|
||||
FirstTokenDocumentOffset = previousState?.FirstTokenDocumentOffset ?? documentOffset,
|
||||
RangeCount = previousState == null ? 1 :
|
||||
previousState.RangeCount + (HasNonEmptyCharacters(documentCodePoints[reference.DocumentId], previousState.DocumentOffset.End, currentStart) ? 1 : 0),
|
||||
TokenCount = (previousState?.TokenCount ?? 0) + 1,
|
||||
PrefixMatchCount = (previousState?.PrefixMatchCount ?? 0) + (isTokenPrefixMatching ? 1 : 0),
|
||||
MatchedTokenLength = (previousState?.MatchedTokenLength ?? 0) + documentOffset.Length *
|
||||
Math.Min(isTokenPrefixMatching ? (double)inputOffset.Length / tokenDefinitions[tokenId].CodePointLength : double.PositiveInfinity, 1),
|
||||
TokenId = tokenId,
|
||||
DocumentOffset = documentOffset,
|
||||
InputOffset = inputOffset,
|
||||
IsTokenPrefixMatching = isTokenPrefixMatching,
|
||||
};
|
||||
nextMatches[currentEnd] = oldResult == null || newResult.CompareTo(oldResult) < 0 ? newResult : oldResult;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Build search results and sort documents
|
||||
return dp[codePoints.Length - 1].Select(entry =>
|
||||
{
|
||||
var (documentId, matches) = entry;
|
||||
var sortedMatches = matches.Values.Select(match =>
|
||||
{
|
||||
var tokens = new List<SearchResultToken>();
|
||||
// Build token list from backtracking
|
||||
var state = match;
|
||||
while (state != null)
|
||||
{
|
||||
tokens.Add(new SearchResultToken
|
||||
{
|
||||
Definition = tokenDefinitions[state.TokenId],
|
||||
DocumentOffset = state.DocumentOffset, InputOffset = state.InputOffset,
|
||||
IsTokenPrefixMatching = state.IsTokenPrefixMatching,
|
||||
});
|
||||
state = state.PreviousState;
|
||||
}
|
||||
tokens.Reverse();
|
||||
return new CandidateResult
|
||||
{
|
||||
Tokens = tokens.ToArray(),
|
||||
PrefixMatchCount = match.PrefixMatchCount,
|
||||
MatchedTokenLength = match.MatchedTokenLength,
|
||||
RangeCount = match.RangeCount,
|
||||
};
|
||||
}).OrderBy(match => match);
|
||||
var bestMatch = sortedMatches.First();
|
||||
var documentText = documents[documentId];
|
||||
var matchRatio = bestMatch.MatchedTokenLength / documentCodePoints[documentId].Length;
|
||||
var matchRatioLevel = (int)Math.Round(matchRatio * 5);
|
||||
return new FinalResult
|
||||
{
|
||||
Result = new SearchResult
|
||||
{
|
||||
DocumentId = documentId,
|
||||
DocumentText = documentText,
|
||||
DocumentCodePoints = documentCodePoints[documentId],
|
||||
Tokens = bestMatch.Tokens,
|
||||
PrefixMatchCount = bestMatch.PrefixMatchCount,
|
||||
RangeCount = bestMatch.RangeCount,
|
||||
MatchRatio = matchRatio,
|
||||
MatchRatioLevel = matchRatioLevel,
|
||||
}
|
||||
};
|
||||
}).OrderBy(result => result).Select(result => result.Result).ToArray();
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
<Project Sdk="Microsoft.NET.Sdk">
|
||||
|
||||
<PropertyGroup>
|
||||
<TargetFramework>netstandard2.0</TargetFramework>
|
||||
<OutputType>Library</OutputType>
|
||||
<RootNamespace>$(ProjectName).Searcher</RootNamespace>
|
||||
<AssemblyName>$(RootNamespace)</AssemblyName>
|
||||
</PropertyGroup>
|
||||
|
||||
<PropertyGroup>
|
||||
<IsPackable>true</IsPackable>
|
||||
<PackageId>$(RootNamespace)</PackageId>
|
||||
</PropertyGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<ProjectReference Include="..\MaigoLabs.NeedLe.Common\MaigoLabs.NeedLe.Common.csproj" />
|
||||
</ItemGroup>
|
||||
|
||||
<ItemGroup>
|
||||
<PackageReference Include="DotNetCampus.LatestCSharpFeatures" PrivateAssets="all" />
|
||||
</ItemGroup>
|
||||
|
||||
</Project>
|
||||
@@ -0,0 +1,37 @@
|
||||
using MaigoLabs.NeedLe.Common.Extensions;
|
||||
using MaigoLabs.NeedLe.Common.Types;
|
||||
|
||||
namespace MaigoLabs.NeedLe.Searcher;
|
||||
|
||||
public class HighlightedTextPart
|
||||
{
|
||||
public required string Text { get; init; }
|
||||
public required bool IsHighlighted { get; init; }
|
||||
}
|
||||
|
||||
public static class SearchResultHighlighter
|
||||
{
|
||||
public static List<HighlightedTextPart> Highlight(SearchResult resultDocument)
|
||||
{
|
||||
var result = new List<HighlightedTextPart>();
|
||||
var previousHighlightEnd = 0;
|
||||
foreach (var token in resultDocument.Tokens)
|
||||
{
|
||||
var notHighlightedText = resultDocument.DocumentCodePoints.Skip(previousHighlightEnd).Take(token.DocumentOffset.Start - previousHighlightEnd).ToUtf32String();
|
||||
if (notHighlightedText.Length > 0) result.Add(new HighlightedTextPart { Text = notHighlightedText, IsHighlighted = false });
|
||||
var highlightEnd = token.IsTokenPrefixMatching && token.Definition.Type == TokenType.Kana
|
||||
? token.DocumentOffset.Start + Math.Max(
|
||||
1,
|
||||
(int)Math.Round(
|
||||
token.DocumentOffset.Length *
|
||||
Math.Min(1, (double)token.InputOffset.Length / token.Definition.CodePointLength)
|
||||
)
|
||||
)
|
||||
: token.DocumentOffset.End;
|
||||
result.Add(new HighlightedTextPart { Text = resultDocument.DocumentCodePoints.Skip(token.DocumentOffset.Start).Take(highlightEnd - token.DocumentOffset.Start).ToUtf32String(), IsHighlighted = true });
|
||||
previousHighlightEnd = highlightEnd;
|
||||
}
|
||||
if (previousHighlightEnd < resultDocument.DocumentCodePoints.Length) result.Add(new HighlightedTextPart { Text = resultDocument.DocumentCodePoints.Skip(previousHighlightEnd).ToUtf32String(), IsHighlighted = false });
|
||||
return result;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,73 @@
|
||||
using MaigoLabs.NeedLe.Common;
|
||||
|
||||
namespace MaigoLabs.NeedLe.Searcher.Trie;
|
||||
|
||||
public class DeserializedTrie
|
||||
{
|
||||
public required TrieNode Root { get; set; }
|
||||
public required Dictionary<int, int[]> TokenCodePoints { get; set; }
|
||||
}
|
||||
|
||||
public static class TrieDeserializer
|
||||
{
|
||||
public static DeserializedTrie Deserialize(int[] data)
|
||||
{
|
||||
var nodes = new List<TrieNode?>();
|
||||
TrieNode GetNode(int id)
|
||||
{
|
||||
if (id > nodes.Count) nodes.AddRange(Enumerable.Repeat<TrieNode?>(null, id - nodes.Count));
|
||||
return nodes[id - 1] ??= new TrieNode { Parent = null, Children = [], TokenIds = [], SubTreeTokenIds = [] };
|
||||
}
|
||||
var currentId = 0;
|
||||
for (var i = 0; i < data.Length; )
|
||||
{
|
||||
var node = GetNode(++currentId);
|
||||
var parentId = data[i++];
|
||||
node.Parent = parentId != 0 ? GetNode(parentId) : null;
|
||||
|
||||
var endOfChildren = i;
|
||||
while (endOfChildren < data.Length && data[endOfChildren] > 0) endOfChildren++;
|
||||
var numberOfChildren = (endOfChildren - i) / 2;
|
||||
for (var j = i; j < i + numberOfChildren; j++)
|
||||
{
|
||||
var codePoint = data[j];
|
||||
var child = GetNode(data[j + numberOfChildren]);
|
||||
node.Children.Add(codePoint, child);
|
||||
}
|
||||
i = endOfChildren;
|
||||
|
||||
if (data[i] == 0) i++; // No token IDs
|
||||
else while (i < data.Length && data[i] < 0) node.TokenIds.Add(-data[i++] - 1);
|
||||
}
|
||||
var root = nodes[0]!;
|
||||
|
||||
// DFS to construct code point paths for each token
|
||||
var tokenCodePoints = new Dictionary<int, int[]>();
|
||||
var currentCodePoints = new List<int>();
|
||||
void DfsCodePoints(TrieNode node)
|
||||
{
|
||||
foreach (var tokenId in node.TokenIds) tokenCodePoints.Add(tokenId, [.. currentCodePoints]);
|
||||
foreach (var (codePoint, child) in node.Children)
|
||||
{
|
||||
if (child.Parent != node) continue; // Skip grafted paths as these are not the canonical representation of the tokens
|
||||
currentCodePoints.Add(codePoint);
|
||||
DfsCodePoints(child);
|
||||
currentCodePoints.RemoveAt(currentCodePoints.Count - 1);
|
||||
}
|
||||
}
|
||||
DfsCodePoints(root);
|
||||
|
||||
// DFS to construct subTreeTokenIds for each node
|
||||
var visitedNodes = new HashSet<TrieNode>();
|
||||
List<int> DfsSubTreeTokenIds(TrieNode node)
|
||||
{
|
||||
if (visitedNodes.Contains(node)) return node.SubTreeTokenIds;
|
||||
visitedNodes.Add(node);
|
||||
node.SubTreeTokenIds = new HashSet<int>(node.TokenIds.Concat(node.Children.Values.SelectMany(DfsSubTreeTokenIds))).ToList();
|
||||
return node.SubTreeTokenIds;
|
||||
};
|
||||
DfsSubTreeTokenIds(root);
|
||||
|
||||
return new DeserializedTrie { Root = root, TokenCodePoints = tokenCodePoints };
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user