You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Zero/ZeroLevel.HNSW/Services/HNSWMap.cs

85 lines
2.3 KiB

using System.Collections.Generic;
using System.IO;
using ZeroLevel.Services.Serialization;
namespace ZeroLevel.HNSW
{
// object -> vector -> vectorId
// HNSW vectorId + vector
// Map object feature - vectorId
public class HNSWMap<TFeature>
: IBinarySerializable
{
private Dictionary<TFeature, int> _map;
private Dictionary<int, TFeature> _reverse_map;
public int this[TFeature feature] => _map.GetValueOrDefault(feature);
public HNSWMap() { }
public HNSWMap(int capacity = -1)
{
if (capacity > 0)
{
_map = new Dictionary<TFeature, int>(capacity);
_reverse_map = new Dictionary<int, TFeature>(capacity);
}
else
{
_map = new Dictionary<TFeature, int>();
_reverse_map = new Dictionary<int, TFeature>();
}
}
public HNSWMap(Stream stream)
{
using (var reader = new MemoryStreamReader(stream))
{
Deserialize(reader);
}
}
public void Append(TFeature feature, int vectorId)
{
_map[feature] = vectorId;
_reverse_map[vectorId] = feature;
}
public IEnumerable<int> ConvertFeaturesToIds(IEnumerable<TFeature> features)
{
int id;
foreach (var feature in features)
{
if (_map.TryGetValue(feature, out id))
{
yield return id;
}
}
}
public IEnumerable<TFeature> ConvertIdsToFeatures(IEnumerable<int> ids)
{
TFeature feature;
foreach (var id in ids)
{
if (_reverse_map.TryGetValue(id, out feature))
{
yield return feature;
}
}
}
public void Deserialize(IBinaryReader reader)
{
this._map = reader.ReadDictionary<TFeature, int>();
this._reverse_map = reader.ReadDictionary<int, TFeature>();
}
public void Serialize(IBinaryWriter writer)
{
writer.WriteDictionary<TFeature, int>(this._map);
writer.WriteDictionary<int, TFeature>(this._reverse_map);
}
}
}

Powered by TurnKey Linux.