Merge pull request #1351 from martindevans/cleanup_mtmd

General MTMD Improvements & Cleanup

Author: martindevans

LLama.Unittest/MtmdWeightsTests.cs

@@ -30,7 +30,7 @@ public MtmdWeightTests()
 
             _mediaMarker = _mtmdParams.MediaMarker ?? throw new InvalidOperationException("MTMD media marker unavailable.");
 
-            _mtmdWeights = MtmdWeights.LoadFromFile(Constants.MtmdMmpPath, _llamaWeights, _mtmdParams);
+            _mtmdWeights = Task.Run(async () => await MtmdWeights.LoadFromFileAsync(Constants.MtmdMmpPath, _llamaWeights, _mtmdParams)).Result;
             _context = _llamaWeights.CreateContext(@params);
         }
 
@@ -53,7 +53,10 @@ private SafeMtmdInputChunks TokenizeWithEmbed(Func<SafeMtmdEmbed> loadEmbed)
                 Assert.True(embed.Nx > 0);
                 Assert.True(embed.Ny > 0);
                 Assert.False(embed.IsAudio);
-                Assert.True(embed.GetDataSpan().Length > 0);
+
+                Assert.True(embed.ByteCount > 0);
+                using var mem = embed.GetData();
+                Assert.True(mem.Data.Length > 0);
 
                 var status = _mtmdWeights.Tokenize(_mediaMarker, addSpecial: true, parseSpecial: true, out var chunks);
                 Assert.Equal(0, status);
@@ -71,6 +74,16 @@ private void AssertChunksEvaluate(SafeMtmdInputChunks chunks)
             Assert.True(nPast > 0);
         }
 
+        [Fact, Trait("Category", "NoCI")]
+        public void BasicPropertyChecks()
+        {
+            Assert.False(_mtmdWeights.SupportsAudio);
+            Assert.True(_mtmdWeights.SupportsVision);
+            Assert.False(_mtmdWeights.UsesMRope);
+            Assert.True(_mtmdWeights.UsesNonCausalAttention);
+            Assert.Equal(-1, _mtmdWeights.AudioBitrate);
+        }
+
         [Fact,Trait("Category", "NoCI")]
         public void EmbedImageAsFileName()
         {
@@ -125,8 +138,8 @@ public void TokenizeProvidesChunkMetadata()
 
             Assert.True(imageChunks > 0);
             Assert.True(totalTokens > 0);
-            Assert.Equal(totalTokens, _mtmdWeights.CountTokens(chunks));
-            Assert.Equal(totalPositions, _mtmdWeights.CountPositions(chunks));
+            Assert.Equal(totalTokens, chunks.CountTokens());
+            Assert.Equal(totalPositions, chunks.CountPositions());
             Assert.True(_mtmdWeights.SupportsVision);
             Assert.False(_mtmdWeights.SupportsAudio);
 

LLama/Batched/Conversation.cs

@@ -101,7 +101,7 @@ public static MtmdChunkSequence Create(SafeMtmdInputChunks chunks, MtmdWeights c
                 }
             }
 
-            var totalPositions = (int)clipModel.CountPositions(chunks);
+            var totalPositions = (int)chunks.CountPositions();
             return new MtmdChunkSequence(chunks, textTokens, totalPositions);
         }
 

LLama/LLamaExecutorBase.cs

@@ -362,7 +362,7 @@ protected Task PreprocessMtmd(string text, InferStateArgs args, bool addBos, boo
                         tokens.Add(token);
                 }
 
-                var totalPositions = (int)ClipModel.CountPositions(chunks);
+                var totalPositions = (int)chunks.CountPositions();
                 var fillerToken = GetFillerToken(marker);
 
                 if (replaceExisting)

LLama/MtmdWeights.cs

@@ -1,36 +1,79 @@
-
 using System;
 using System.Threading;
 using System.Threading.Tasks;
+using LLama.Exceptions;
 using LLama.Native;
 
 namespace LLama;
 
 /// <summary>
 /// Lightweight wrapper around the MTMD native context and its helpers.
 /// </summary>
-public sealed class MtmdWeights : IDisposable
+public sealed class MtmdWeights
+    : IDisposable
 {
+    /// <summary>
+    /// The native handle, which is used in the native APIs
+    /// </summary>
+    /// <remarks>Be careful how you use this!</remarks>
     public SafeMtmdModelHandle NativeHandle { get; }
 
     private MtmdWeights(SafeMtmdModelHandle handle)
     {
         NativeHandle = handle ?? throw new ArgumentNullException(nameof(handle));
     }
 
+    /// <summary>
+    /// Load weights into memory
+    /// </summary>
+    /// <param name="mmProject">Path to the mmproj file</param>
+    /// <param name="textModel">The text model</param>
+    /// <param name="mtmdCtxParams">Parameters for MTMD context creation</param>
+    /// <returns></returns>
     public static MtmdWeights LoadFromFile(string mmProject, LLamaWeights textModel, MtmdContextParams mtmdCtxParams)
     {
-        if (mmProject == null) throw new ArgumentNullException(nameof(mmProject));
-        if (textModel == null) throw new ArgumentNullException(nameof(textModel));
-        if (mtmdCtxParams == null) throw new ArgumentNullException(nameof(mtmdCtxParams));
-
-        var handle = SafeMtmdModelHandle.LoadFromFile(mmProject, textModel, mtmdCtxParams);
-        return new MtmdWeights(handle);
+        return new MtmdWeights(SafeMtmdModelHandle.LoadFromFile(
+            mmProject ?? throw new ArgumentNullException(nameof(mmProject)),
+            textModel ?? throw new ArgumentNullException(nameof(textModel)),
+            mtmdCtxParams ?? throw new ArgumentNullException(nameof(mtmdCtxParams))
+        ));
     }
 
-    public static Task<MtmdWeights> LoadFromFileAsync(string mmProject, LLamaWeights textModel, MtmdContextParams mtmdCtxParams, CancellationToken token = default)
+    /// <summary>
+    /// Load weights into memory
+    /// </summary>
+    /// <param name="mmProject">Path to the mmproj file</param>
+    /// <param name="textModel">The text model</param>
+    /// <param name="mtmdCtxParams">Parameters for MTMD context creation</param>
+    /// <param name="token"></param>
+    /// <returns></returns>
+    public static async Task<MtmdWeights> LoadFromFileAsync(string mmProject, LLamaWeights textModel, MtmdContextParams mtmdCtxParams, CancellationToken token = default)
     {
-        return Task.Run(() => LoadFromFile(mmProject, textModel, mtmdCtxParams), token);
+        if (mmProject == null)
+            throw new ArgumentNullException(nameof(mmProject));
+        if (textModel == null)
+            throw new ArgumentNullException(nameof(textModel));
+        if (mtmdCtxParams == null)
+            throw new ArgumentNullException(nameof(mtmdCtxParams));
+
+        var model = await Task.Run(() =>
+        {
+            try
+            {
+                // Load the model
+                return LoadFromFile(mmProject, textModel, mtmdCtxParams);
+            }
+            catch (LoadWeightsFailedException)
+            {
+                // Convert a LoadWeightsFailedException into a cancellation exception if possible.
+                token.ThrowIfCancellationRequested();
+
+                // Ok the weights failed to load for some reason other than cancellation.
+                throw;
+            }
+        }, token);
+
+        return model;
     }
 
     /// <summary>
@@ -73,15 +116,31 @@ public int EvaluateChunk(IntPtr chunkPtr, SafeLLamaContextHandle llamaContext, r
     public int DecodeImageChunk(IntPtr chunkPtr, SafeLLamaContextHandle llamaContext, IntPtr encodedEmbeddings, ref int nPast, int seqId, int nBatch)
         => NativeHandle.DecodeImageChunk(chunkPtr, llamaContext, encodedEmbeddings, ref nPast, seqId, nBatch);
 
-    public ulong CountTokens(SafeMtmdInputChunks chunks) => NativeHandle.CountTokens(chunks);
-
-    public long CountPositions(SafeMtmdInputChunks chunks) => NativeHandle.CountPositions(chunks);
-
+    /// <summary>
+    /// Indicates whether the model supports vision inputs.
+    /// </summary>
     public bool SupportsVision => NativeHandle.SupportVision();
+
+    /// <summary>
+    /// Indicates whether the model supports audio inputs.
+    /// </summary>
     public bool SupportsAudio => NativeHandle.SupportAudio();
+
+    /// <summary>
+    /// Indicates whether the model decodes using the non-causal path.
+    /// </summary>
     public bool UsesNonCausalAttention => NativeHandle.DecodeUseNonCausal();
+
+    /// <summary>
+    /// Indicates whether the model decodes using multi-scale RoPE.
+    /// </summary>
     public bool UsesMRope => NativeHandle.DecodeUseMRope();
+
+    /// <summary>
+    /// Gets the audio bitrate advertised by the model.
+    /// </summary>
     public int AudioBitrate => NativeHandle.GetAudioBitrate();
 
+    /// <inheritdoc />
     public void Dispose() => NativeHandle.Dispose();
 }