Search.cs

Go to the documentation of this file.

 
using System;
using System.Collections.Generic;
using System.IO;
using System.IO.Compression;
using System.Linq;
using System.Runtime.Serialization.Formatters.Binary;
using System.Threading.Tasks;
using UnityEditor;
using UnityEngine;
 
namespace LLMUnity
{
    [DefaultExecutionOrder(-2)]
    public abstract class Searchable : MonoBehaviour
    {
        public abstract string Get(int key);
 
        public abstract Task<int> Add(string inputString, string group = "");
 
        public abstract int Remove(string inputString, string group = "");
 
        public abstract void Remove(int key);
 
        public abstract int Count();
 
        public abstract int Count(string group);
 
        public abstract void Clear();
 
        public abstract Task<int> IncrementalSearch(string queryString, string group = "");
 
        public abstract ValueTuple<int[], float[], bool> IncrementalFetchKeys(int fetchKey, int k);
 
        public abstract void IncrementalSearchComplete(int fetchKey);
 
        public async Task<(string[], float[])> Search(string queryString, int k, string group = "")
        {
            int fetchKey = await IncrementalSearch(queryString, group);
            (string[] phrases, float[] distances, bool completed) = IncrementalFetch(fetchKey, k);
            if (!completed) IncrementalSearchComplete(fetchKey);
            return (phrases, distances);
        }
 
        public virtual ValueTuple<string[], float[], bool> IncrementalFetch(int fetchKey, int k)
        {
            (int[] resultKeys, float[] distances, bool completed) = IncrementalFetchKeys(fetchKey, k);
            string[] results = new string[resultKeys.Length];
            for (int i = 0; i < resultKeys.Length; i++) results[i] = Get(resultKeys[i]);
            return (results, distances, completed);
        }
 
        public void Save(string filePath)
        {
            try
            {
                string path = LLMUnitySetup.GetAssetPath(filePath);
                ArchiveSaver.Save(path, Save);
            }
            catch (Exception e)
            {
                LLMUnitySetup.LogError($"File {filePath} could not be saved due to {e.GetType()}: {e.Message}");
            }
        }
 
        public async Task<bool> Load(string filePath)
        {
            try
            {
                await LLMUnitySetup.AndroidExtractAsset(filePath, true);
                string path = LLMUnitySetup.GetAssetPath(filePath);
                if (!File.Exists(path)) return false;
                ArchiveSaver.Load(path, Load);
            }
            catch (Exception e)
            {
                LLMUnitySetup.LogError($"File {filePath} could not be loaded due to {e.GetType()}: {e.Message}");
                return false;
            }
            return true;
        }
 
        public abstract void Save(ZipArchive archive);
        public abstract void Load(ZipArchive archive);
        public virtual string GetSavePath(string name)
        {
            return Path.Combine(GetType().Name, name);
        }
 
        public virtual void UpdateGameObjects() {}
 
        protected T ConstructComponent<T>(Type type, Action<T, T> copyAction = null) where T : Component
        {
            T Construct(Type type)
            {
                if (type == null) return null;
                T newComponent = (T)gameObject.AddComponent(type);
                if (newComponent is Searchable searchable) searchable.UpdateGameObjects();
                return newComponent;
            }
 
            T component = (T)gameObject.GetComponent(typeof(T));
            T newComponent;
            if (component == null)
            {
                newComponent = Construct(type);
            }
            else
            {
                if (component.GetType() == type)
                {
                    newComponent = component;
                }
                else
                {
                    newComponent = Construct(type);
                    if (type != null) copyAction?.Invoke(component, newComponent);
#if UNITY_EDITOR
                    DestroyImmediate(component);
#else
                    Destroy(component);
#endif
                }
            }
            return newComponent;
        }
 
        public virtual void Awake()
        {
            UpdateGameObjects();
        }
 
#if UNITY_EDITOR
        public virtual void Reset()
        {
            if (!Application.isPlaying) EditorApplication.update += UpdateGameObjects;
        }
 
        public virtual void OnDestroy()
        {
            if (!Application.isPlaying) EditorApplication.update -= UpdateGameObjects;
        }
 
#endif
    }
 
    public abstract class SearchMethod : Searchable
    {
        public LLMEmbedder llmEmbedder;
 
        protected int nextKey = 0;
        protected int nextIncrementalSearchKey = 0;
        protected SortedDictionary<int, string> data = new SortedDictionary<int, string>();
        protected SortedDictionary<string, List<int>> dataSplits = new SortedDictionary<string, List<int>>();
 
        protected LLM llm;
 
        protected abstract void AddInternal(int key, float[] embedding);
        protected abstract void RemoveInternal(int key);
        protected abstract void ClearInternal();
        protected abstract void SaveInternal(ZipArchive archive);
        protected abstract void LoadInternal(ZipArchive archive);
 
        public void SetLLM(LLM llm)
        {
            this.llm = llm;
            if (llmEmbedder != null) llmEmbedder.llm = llm;
        }
 
        public async Task<(string[], float[])> SearchFromList(string query, string[] searchList)
        {
            float[] embedding = await Encode(query);
            float[][] embeddingsList = new float[searchList.Length][];
            for (int i = 0; i < searchList.Length; i++) embeddingsList[i] = await Encode(searchList[i]);
 
            float[] unsortedDistances = InverseDotProduct(embedding, embeddingsList);
            List<(string, float)> sortedLists = searchList.Zip(unsortedDistances, (first, second) => (first, second))
                .OrderBy(item => item.Item2)
                .ToList();
 
            string[] results = new string[sortedLists.Count];
            float[] distances = new float[sortedLists.Count];
            for (int i = 0; i < sortedLists.Count; i++)
            {
                results[i] = sortedLists[i].Item1;
                distances[i] = sortedLists[i].Item2;
            }
            return (results.ToArray(), distances.ToArray());
        }
 
        public static float DotProduct(float[] vector1, float[] vector2)
        {
            if (vector1 == null || vector2 == null) throw new ArgumentNullException("Vectors cannot be null");
            if (vector1.Length != vector2.Length) throw new ArgumentException("Vector lengths must be equal for dot product calculation");
            float result = 0;
            for (int i = 0; i < vector1.Length; i++)
            {
                result += vector1[i] * vector2[i];
            }
            return result;
        }
 
        public static float InverseDotProduct(float[] vector1, float[] vector2)
        {
            return 1 - DotProduct(vector1, vector2);
        }
 
        public static float[] InverseDotProduct(float[] vector1, float[][] vector2)
        {
            float[] results = new float[vector2.Length];
            for (int i = 0; i < vector2.Length; i++)
            {
                results[i] = InverseDotProduct(vector1, vector2[i]);
            }
            return results;
        }
 
        public virtual async Task<float[]> Encode(string inputString)
        {
            return (await llmEmbedder.Embeddings(inputString)).ToArray();
        }
 
        public virtual async Task<List<int>> Tokenize(string query, Callback<List<int>> callback = null)
        {
            return await llmEmbedder.Tokenize(query, callback);
        }
 
        public async Task<string> Detokenize(List<int> tokens, Callback<string> callback = null)
        {
            return await llmEmbedder.Detokenize(tokens, callback);
        }
 
        public override string Get(int key)
        {
            if (data.TryGetValue(key, out string result)) return result;
            return null;
        }
 
        public override async Task<int> Add(string inputString, string group = "")
        {
            int key = nextKey++;
            AddInternal(key, await Encode(inputString));
 
            data[key] = inputString;
            if (!dataSplits.ContainsKey(group)) dataSplits[group] = new List<int>(){key};
            else dataSplits[group].Add(key);
            return key;
        }
 
        public override void Clear()
        {
            data.Clear();
            dataSplits.Clear();
            ClearInternal();
            nextKey = 0;
            nextIncrementalSearchKey = 0;
        }
 
        protected bool RemoveEntry(int key)
        {
            bool removed = data.Remove(key);
            if (removed) RemoveInternal(key);
            return removed;
        }
 
        public override void Remove(int key)
        {
            if (RemoveEntry(key))
            {
                foreach (var dataSplit in dataSplits.Values) dataSplit.Remove(key);
            }
        }
 
        public override int Remove(string inputString, string group = "")
        {
            if (!dataSplits.TryGetValue(group, out List<int> dataSplit)) return 0;
            List<int> removeIds = new List<int>();
            foreach (int key in dataSplit)
            {
                if (Get(key) == inputString) removeIds.Add(key);
            }
            foreach (int key in removeIds)
            {
                if (RemoveEntry(key)) dataSplit.Remove(key);
            }
            return removeIds.Count;
        }
 
        public override int Count()
        {
            return data.Count;
        }
 
        public override int Count(string group)
        {
            if (!dataSplits.TryGetValue(group, out List<int> dataSplit)) return 0;
            return dataSplit.Count;
        }
 
        public override async Task<int> IncrementalSearch(string queryString, string group = "")
        {
            return IncrementalSearch(await Encode(queryString), group);
        }
 
        public override void Save(ZipArchive archive)
        {
            ArchiveSaver.Save(archive, data, GetSavePath("data"));
            ArchiveSaver.Save(archive, dataSplits, GetSavePath("dataSplits"));
            ArchiveSaver.Save(archive, nextKey, GetSavePath("nextKey"));
            ArchiveSaver.Save(archive, nextIncrementalSearchKey, GetSavePath("nextIncrementalSearchKey"));
            SaveInternal(archive);
        }
 
        public override void Load(ZipArchive archive)
        {
            data = ArchiveSaver.Load<SortedDictionary<int, string>>(archive, GetSavePath("data"));
            dataSplits = ArchiveSaver.Load<SortedDictionary<string, List<int>>>(archive, GetSavePath("dataSplits"));
            nextKey = ArchiveSaver.Load<int>(archive, GetSavePath("nextKey"));
            nextIncrementalSearchKey = ArchiveSaver.Load<int>(archive, GetSavePath("nextIncrementalSearchKey"));
            LoadInternal(archive);
        }
 
        public override void UpdateGameObjects()
        {
            if (this == null || llmEmbedder != null) return;
            llmEmbedder = ConstructComponent<LLMEmbedder>(typeof(LLMEmbedder), (previous, current) => current.llm = previous.llm);
        }
 
        public abstract int IncrementalSearch(float[] embedding, string group = "");
    }
 
    public abstract class SearchPlugin : Searchable
    {
        protected SearchMethod search;
 
        public void SetSearch(SearchMethod search)
        {
            this.search = search;
        }
 
        protected abstract void SaveInternal(ZipArchive archive);
        protected abstract void LoadInternal(ZipArchive archive);
 
        public override void Save(ZipArchive archive)
        {
            search.Save(archive);
            SaveInternal(archive);
        }
 
        public override void Load(ZipArchive archive)
        {
            search.Load(archive);
            LoadInternal(archive);
        }
 
    }
 
    public class ArchiveSaver
    {
        public delegate void ArchiveSaverCallback(ZipArchive archive);
 
        public static void Save(string filePath, ArchiveSaverCallback callback)
        {
            using (FileStream stream = new FileStream(filePath, FileMode.Create))
            using (ZipArchive archive = new ZipArchive(stream, ZipArchiveMode.Create))
            {
                callback(archive);
            }
        }
 
        public static void Load(string filePath, ArchiveSaverCallback callback)
        {
            using (FileStream stream = new FileStream(filePath, FileMode.Open))
            using (ZipArchive archive = new ZipArchive(stream, ZipArchiveMode.Read))
            {
                callback(archive);
            }
        }
 
        public static void Save(ZipArchive archive, object saveObject, string name)
        {
            ZipArchiveEntry mainEntry = archive.CreateEntry(name);
            using (Stream entryStream = mainEntry.Open())
            {
                BinaryFormatter formatter = new BinaryFormatter();
                formatter.Serialize(entryStream, saveObject);
            }
        }
 
        public static T Load<T>(ZipArchive archive, string name)
        {
            ZipArchiveEntry baseEntry = archive.GetEntry(name);
            if (baseEntry == null) throw new Exception($"No entry with name {name} was found");
            using (Stream entryStream = baseEntry.Open())
            {
                BinaryFormatter formatter = new BinaryFormatter();
                return (T)formatter.Deserialize(entryStream);
            }
        }
    }
}