Implement TopProbabilityNum and StopSequences for ollamaChat

Author: ccreutzi

+llms/+internal/callOllamaChatAPI.m

@@ -25,31 +25,16 @@
 %
 %   Example
 %
+%   model = "mistral";
+%
 %   % Create messages struct
 %   messages = {struct("role", "system",...
 %       "content", "You are a helpful assistant");
 %       struct("role", "user", ...
 %       "content", "What is the edit distance between hi and hello?")};
 %
-%   % Create functions struct
-%   functions = {struct("name", "editDistance", ...
-%       "description", "Find edit distance between two strings or documents.", ...
-%       "parameters", struct( ...
-%       "type", "object", ...
-%       "properties", struct(...
-%           "str1", struct(...
-%               "description", "Source string.", ...
-%               "type", "string"),...
-%           "str2", struct(...
-%               "description", "Target string.", ...
-%               "type", "string")),...
-%       "required", ["str1", "str2"]))};
-%
-%   % Define your API key
-%   apiKey = "your-api-key-here"
-%
 %   % Send a request
-%   [text, message] = llms.internal.callOpenAIChatAPI(messages, functions, ApiKey=apiKey)
+%   [text, message] = llms.internal.callOllamaChatAPI(model, messages)
 
 %   Copyright 2023-2024 The MathWorks, Inc.
 
@@ -58,6 +43,7 @@
     messages
     nvp.Temperature = 1
     nvp.TopProbabilityMass = 1
+    nvp.TopProbabilityNum = Inf
     nvp.NumCompletions = 1
     nvp.StopSequences = []
     nvp.MaxNumTokens = inf
@@ -71,6 +57,12 @@
 
 URL = "http://localhost:11434/api/chat"; % TODO: model parameter
 
+% The JSON for StopSequences must have an array, and cannot say "stop": "foo".
+% The easiest way to ensure that is to never pass in a scalar …
+if isscalar(nvp.StopSequences)
+    nvp.StopSequences = [nvp.StopSequences, nvp.StopSequences];
+end
+
 parameters = buildParametersCall(model, messages, nvp);
 
 [response, streamedText] = llms.internal.sendRequest(parameters,[],URL,nvp.TimeOut,nvp.StreamFun);
@@ -123,6 +115,7 @@
 dict = dictionary();
 dict("Temperature") = "temperature";
 dict("TopProbabilityMass") = "top_p";
+dict("TopProbabilityNum") = "top_k";
 dict("NumCompletions") = "n";
 dict("StopSequences") = "stop";
 dict("MaxNumTokens") = "num_predict";

ollamaChat.m

@@ -1,4 +1,4 @@
-classdef(Sealed) ollamaChat < llms.internal.textGenerator
+classdef (Sealed) ollamaChat < llms.internal.textGenerator
 %ollamaChat Chat completion API from Azure.
 %
 %   CHAT = ollamaChat(modelName) creates an ollamaChat object for the given model.
@@ -68,8 +68,9 @@
 
 % Copyright 2024 The MathWorks, Inc.
 
-    properties(SetAccess=private)
+    properties
         Model     (1,1) string
+        TopProbabilityNum (1,1) {mustBeReal,mustBePositive} = Inf
     end
 
     methods
@@ -79,6 +80,7 @@
                 systemPrompt                       {llms.utils.mustBeTextOrEmpty} = []
                 nvp.Temperature                    {llms.utils.mustBeValidTemperature} = 1
                 nvp.TopProbabilityMass             {llms.utils.mustBeValidTopP} = 1
+                nvp.TopProbabilityNum        (1,1) {mustBeReal,mustBePositive} = Inf
                 nvp.StopSequences                  {llms.utils.mustBeValidStop} = {}
                 nvp.ResponseFormat           (1,1) string {mustBeMember(nvp.ResponseFormat,["text","json"])} = "text"
                 nvp.TimeOut                  (1,1) {mustBeReal,mustBePositive} = 10
@@ -102,6 +104,7 @@
             this.ResponseFormat = nvp.ResponseFormat;
             this.Temperature = nvp.Temperature;
             this.TopProbabilityMass = nvp.TopProbabilityMass;
+            this.TopProbabilityNum = nvp.TopProbabilityNum;
             this.StopSequences = nvp.StopSequences;
             this.TimeOut = nvp.TimeOut;
         end
@@ -152,6 +155,7 @@
             [text, message, response] = llms.internal.callOllamaChatAPI(...
                 this.Model, messagesStruct, ...
                 Temperature=this.Temperature, ...
+                TopProbabilityMass=this.TopProbabilityMass, TopProbabilityNum=this.TopProbabilityNum,...
                 NumCompletions=nvp.NumCompletions,...
                 StopSequences=this.StopSequences, MaxNumTokens=nvp.MaxNumTokens, ...
                 ResponseFormat=this.ResponseFormat,Seed=nvp.Seed, ...

tests/tollamaChat.m

@@ -7,6 +7,7 @@
         InvalidConstructorInput = iGetInvalidConstructorInput;
         InvalidGenerateInput = iGetInvalidGenerateInput;
         InvalidValuesSetters = iGetInvalidValuesSetters;
+        ValidValuesSetters = iGetValidValuesSetters;
     end
 
     methods(Test)
@@ -37,6 +38,26 @@ function doGenerate(testCase)
             testCase.verifyGreaterThan(strlength(response),0);
         end
 
+        function extremeTopK(testCase)
+            % setting top-k to k=1 leaves no random choice,
+            % so we expect to get a fixed response.
+            chat = ollamaChat("mistral",TopProbabilityNum=1);
+            prompt = "Top-k sampling with k=1 returns a definite answer.";
+            response1 = generate(chat,prompt);
+            response2 = generate(chat,prompt);
+            testCase.verifyEqual(response1,response2);
+        end
+
+        function stopSequences(testCase)
+            chat = ollamaChat("mistral",TopProbabilityNum=1);
+            prompt = "Top-k sampling with k=1 returns a definite answer.";
+            response1 = generate(chat,prompt);
+            chat.StopSequences = "1";
+            response2 = generate(chat,prompt);
+
+            testCase.verifyEqual(response2, extractBefore(response1,"1"));
+        end
+
         %% Test is currently unreliable, reasons unclear
         % function verySmallTimeOutErrors(testCase)
         %     chat = ollamaChat("mistral", TimeOut=1e-10);
@@ -60,6 +81,15 @@ function assignValueToProperty(property, value)
 
             testCase.verifyError(@() assignValueToProperty(InvalidValuesSetters.Property,InvalidValuesSetters.Value), InvalidValuesSetters.Error);
         end
+
+        function validSetters(testCase, ValidValuesSetters)
+            chat = ollamaChat("mistral");
+            function assignValueToProperty(property, value)
+                chat.(property) = value;
+            end
+
+            testCase.verifyWarningFree(@() assignValueToProperty(ValidValuesSetters.Property,ValidValuesSetters.Value));
+        end
     end
 end
 
@@ -119,17 +149,19 @@ function assignValueToProperty(property, value)
     "EmptyStopSequences", struct( ...
         "Property", "StopSequences", ...
         "Value", "", ...
-        "Error", "MATLAB:validators:mustBeNonzeroLengthText"), ...
-    ...
-    "WrongSizeStopSequences", struct( ...
-        "Property", "StopSequences", ...
-        "Value", ["1" "2" "3" "4" "5"], ...
-        "Error", "llms:stopSequencesMustHaveMax4Elements"), ...
-    ...
-    "InvalidPresencePenalty", struct( ...
-        "Property", "PresencePenalty", ...
-        "Value", "2", ...
-        "Error", "MATLAB:invalidType"));
+        "Error", "MATLAB:validators:mustBeNonzeroLengthText"));
+end
+
+function validSetters = iGetValidValuesSetters
+validSetters = struct(...
+    "SmallTopNum", struct( ...
+        "Property", "TopProbabilityNum", ...
+        "Value", 2));
+    % Currently disabled because it requires some code reorganization
+    % and we have higher priorities ...
+    % "ManyStopSequences", struct( ...
+    %     "Property", "StopSequences", ...
+    %     "Value", ["1" "2" "3" "4" "5"]));
 end
 
 function invalidConstructorInput = iGetInvalidConstructorInput
@@ -196,7 +228,7 @@ function assignValueToProperty(property, value)
     ...
     "TopProbabilityMassTooSmall",struct( ...
         "Input",{{ "TopProbabilityMass" -20 }},...
-        "Error","MATLAB:expectedNonnegative"),...
+        "Error","MATLAB:expectedNonnegative"),...I
     ...
     "WrongTypeStopSequences",struct( ...
         "Input",{{ "StopSequences" 123}},...
@@ -208,11 +240,7 @@ function assignValueToProperty(property, value)
     ...
     "EmptyStopSequences",struct( ...
         "Input",{{ "StopSequences" ""}},...
-        "Error","MATLAB:validators:mustBeNonzeroLengthText"),...
-    ...
-    "WrongSizeStopSequences",struct( ...
-        "Input",{{ "StopSequences" ["1" "2" "3" "4" "5"]}},...
-        "Error","llms:stopSequencesMustHaveMax4Elements"));
+        "Error","MATLAB:validators:mustBeNonzeroLengthText"));
 end
 
 function invalidGenerateInput = iGetInvalidGenerateInput