Model File Manager

LLama.Unittest/Constants.cs

@@ -4,6 +4,7 @@ namespace LLama.Unittest
 {
     internal static class Constants
     {
+        public static readonly string ModelDirectory = "Models";
         public static readonly string GenerativeModelPath = "Models/llama-2-7b-chat.Q3_K_S.gguf";
         public static readonly string EmbeddingModelPath = "Models/all-MiniLM-L12-v2.Q8_0.gguf";
 

LLama.Unittest/LLama.Unittest.csproj

@@ -32,8 +32,6 @@
     <DownloadFile SourceUrl="https://huggingface.co/cjpais/llava-1.6-mistral-7b-gguf/resolve/main/llava-v1.6-mistral-7b.Q3_K_XS.gguf" DestinationFolder="Models" DestinationFileName="llava-v1.6-mistral-7b.Q3_K_XS.gguf" SkipUnchangedFiles="true"></DownloadFile>
     <DownloadFile SourceUrl="https://huggingface.co/cjpais/llava-1.6-mistral-7b-gguf/resolve/main/mmproj-model-f16.gguf" DestinationFolder="Models" DestinationFileName="mmproj-model-f16.gguf" SkipUnchangedFiles="true"></DownloadFile>
     <DownloadFile SourceUrl="https://huggingface.co/leliuga/all-MiniLM-L12-v2-GGUF/resolve/main/all-MiniLM-L12-v2.Q8_0.gguf" DestinationFolder="Models" DestinationFileName="all-MiniLM-L12-v2.Q8_0.gguf" SkipUnchangedFiles="true"></DownloadFile>
-
-
   </Target>
 
   <ItemGroup>

LLama.Unittest/Model/FileSystemModelRepoTests.cs

@@ -0,0 +1,104 @@
+using LLama.Model;
+
+namespace LLama.Unittest.Model;
+
+public class FileSystemModelRepoTests
+{
+    private readonly FileSystemModelRepo TestableRepo;
+
+    public FileSystemModelRepoTests()
+    {
+        TestableRepo = new([Constants.ModelDirectory]);
+    }
+
+    [Fact]
+    public void ModelDirectories_IsCorrect()
+    {
+        var dirs = TestableRepo.ListSources();
+        Assert.Single(dirs);
+
+        var expected = dirs.First()!.Contains(Constants.ModelDirectory);
+        Assert.True(expected);
+    }
+
+    [Fact]
+    public void AddDirectory_DoesntDuplicate()
+    {
+        for (var i = 0; i < 10; i++)
+        {
+            TestableRepo.AddSource(Constants.ModelDirectory);
+            TestableRepo.AddSource(Path.GetFullPath(Constants.ModelDirectory));
+
+            var dirs = TestableRepo.ListSources();
+            Assert.Single(dirs);
+            var expected = dirs.First()!.Contains(Constants.ModelDirectory);
+            Assert.True(expected);
+        }
+    }
+
+    [Fact]
+    public void RemoveDirectory()
+    {
+        var dirs = TestableRepo.ListSources();
+        Assert.Single(dirs);
+        var expected = dirs.First()!.Contains(Constants.ModelDirectory);
+        Assert.True(expected);
+
+        Assert.True(TestableRepo.RemoveSource(Constants.ModelDirectory));
+        Assert.Empty(TestableRepo.ListSources());
+        Assert.Empty(TestableRepo.GetAvailableModels());
+    }
+
+    [Fact]
+    public void RemoveDirectory_DoesNotExist()
+    {
+        var dirs = TestableRepo.ListSources();
+        Assert.Single(dirs);
+        var expected = dirs.First()!.Contains(Constants.ModelDirectory);
+        Assert.True(expected);
+
+        Assert.False(TestableRepo.RemoveSource("foo/boo/bar"));
+        Assert.Single(dirs);
+    }
+
+    [Fact]
+    public void RemoveAllDirectories()
+    {
+        var dirs = TestableRepo.ListSources();
+        Assert.Single(dirs);
+        var expected = dirs.First()!.Contains(Constants.ModelDirectory);
+        Assert.True(expected);
+
+        TestableRepo.RemoveAllSources();
+        Assert.Empty(TestableRepo.ListSources());
+        Assert.Empty(TestableRepo.GetAvailableModels());
+    }
+
+    [Fact]
+    public void ModelFiles_IsCorrect()
+    {
+        var files = TestableRepo.GetAvailableModels();
+        Assert.Equal(4, files.Count());
+    }
+
+    [Fact]
+    public void GetAvailableModelsFromDirectory()
+    {
+        var files = TestableRepo.GetAvailableModelsFromSource(Constants.ModelDirectory);
+        Assert.Equal(4, files.Count());
+
+        files = TestableRepo.GetAvailableModels();
+        Assert.Equal(4, files.Count());
+    }
+
+    [Fact]
+    public void TryGetModelFileMetadata_WhenExists()
+    {
+        var expectedFile = TestableRepo.GetAvailableModels().First();
+        var found = TestableRepo.TryGetModelFileMetadata(expectedFile.ModelFileUri, out var foundData);
+
+        Assert.True(found);
+        Assert.Equal(expectedFile.ModelFileUri, foundData.ModelFileUri);
+    }
+
+}

LLama.Unittest/Model/ModelCacheTests.cs

@@ -0,0 +1,113 @@
+using LLama.Common;
+using LLama.Model;
+
+namespace LLama.Unittest.Model;
+
+public class ModelManagerTests
+{
+    private readonly IModelSourceRepo _testRepo = new FileSystemModelRepo([Constants.ModelDirectory]);
+
+    private readonly ModelCache TestableModelManager;
+
+    public ModelManagerTests()
+    {
+        TestableModelManager = new();
+    }
+
+    [Fact]
+    public async void LoadModel_DisposesOnUnload()
+    {
+        var modelToLoad = _testRepo.GetAvailableModels()
+            .First(f => f.ModelFileName.Contains("llama-2-7b"));
+
+        var model = await TestableModelManager.LoadModelAsync(modelToLoad);
+        Assert.NotNull(model);
+
+        // unloaded and disposed`
+        Assert.True(TestableModelManager.UnloadModel(model.ModelName));
+        Assert.Throws<ObjectDisposedException>(() =>
+        {
+            _ = model.CreateContext(new ModelParams(modelToLoad.ModelFileUri));
+        });
+
+        // wont unload and already
+        Assert.False(TestableModelManager.UnloadModel(model.ModelName));
+        Assert.Throws<ObjectDisposedException>(() =>
+        {
+            _ = model.CreateContext(new ModelParams(modelToLoad.ModelFileUri));
+        });
+    }
+
+    [Fact]
+    public async void LoadModel_LoadsAndCaches()
+    {
+        var modelToLoad = _testRepo.GetAvailableModels()
+            .First(f => f.ModelFileName.Contains("llama-2-7b"));
+
+        // Create Model -- Ref 1
+        var model = await TestableModelManager.LoadModelAsync(modelToLoad);
+        Assert.NotNull(model);
+
+        // clone it -- Ref 2
+        var isCachedAndCloned = TestableModelManager.TryCloneLoadedModel(model.ModelName, out var cachedModel);
+        Assert.True(isCachedAndCloned);
+        Assert.NotNull(cachedModel);
+
+        cachedModel.Dispose(); //-- ref 1
+        Assert.True(TestableModelManager.UnloadModel(model.ModelName));
+
+        // unloaded and disposed` -- ref 2
+        Assert.True(TestableModelManager.UnloadModel(model.ModelName));
+
+        Assert.False(TestableModelManager.UnloadModel(model.ModelName));
+        Assert.Throws<ObjectDisposedException>(() =>
+        {
+            _ = model.CreateContext(new ModelParams(modelToLoad.ModelFileUri));
+        });
+    }
+
+    [Fact]
+    public async void LoadModel_AlreadyLoaded_ReturnsFromCache()
+    {
+        var modelToLoad = _testRepo.GetAvailableModels()
+            .First(f => f.ModelFileName.Contains("llama-2-7b"));
+
+        for (var i = 0; i < 5; i++)
+        {
+            var model = await TestableModelManager.LoadModelAsync(modelToLoad);
+            Assert.NotNull(model);
+            Assert.Equal("LLaMA v2", model.ModelName);
+            var isLoaded = TestableModelManager.TryCloneLoadedModel(model.ModelName, out var cachedModel);
+            Assert.True(isLoaded);
+            Assert.NotNull(cachedModel);
+            Assert.Equal("LLaMA v2", cachedModel.ModelName);
+        }
+    }
+
+    [Fact]
+    public async void TryGetLoadedModel_AlreadyDisposed_ReturnsFalse()
+    {
+        var modelToLoad = _testRepo.GetAvailableModels()
+            .First(f => f.ModelFileName.Contains("llama-2-7b"));
+
+        using (var model = await TestableModelManager.LoadModelAsync(modelToLoad))
+        {
+            Assert.NotNull(model);
+            Assert.Equal(model.ModelName, model.ModelName);
+            var isLoaded = TestableModelManager.TryCloneLoadedModel(model.ModelName, out var cachedModel);
+            Assert.True(isLoaded);
+            Assert.NotNull(cachedModel);
+            Assert.Equal(model.ModelName, cachedModel.ModelName);
+
+            // unload from the last check
+            Assert.True(TestableModelManager.UnloadModel(model.ModelName));
+
+        } // end scope, dispose is called on the model but since we have the model cache it should stick around until unloaded
+        Assert.True(TestableModelManager.UnloadModel("LLaMA v2"));
+
+        // Model is still loaded due to cache
+        var isDisposedLoaded = TestableModelManager.TryCloneLoadedModel("LLaMA v2", out var disposedModel);
+        Assert.False(isDisposedLoaded);
+        Assert.Null(disposedModel);
+    }
+}