feat: add transforms for chat session.

LLama.Examples/NewVersion/LoadAndSaveState.cs

@@ -12,11 +12,9 @@ namespace LLama.Examples.NewVersion
     {
         InteractiveExecutor _executor;
         string _prompt;
-        string _modelPath;
         public SaveAndLoadState(string modelPath, string prompt)
         {
             _prompt = prompt;
-            _modelPath = modelPath;
             _executor = new InteractiveExecutor(new LLamaModel(new ModelParams(modelPath: modelPath)));
             foreach (var text in _executor.Infer(_prompt, new InferenceParams() { Temperature = 0.6f, AntiPrompts = new List<string> { "user:" } }))
             {

LLama.Examples/OldVersion/ChatSession.cs

@@ -5,7 +5,7 @@ using System.Text;
 using System.Threading.Tasks;
 using LLama.OldVersion;
 
-namespace LLama.Examples
+namespace LLama.Examples.Old
 {
     public class ChatSession
     {

LLama.Examples/OldVersion/GetEmbeddings.cs

@@ -5,7 +5,7 @@ using System.Text;
 using System.Threading.Tasks;
 using LLama.OldVersion;
 
-namespace LLama.Examples
+namespace LLama.Examples.Old
 {
     public class GetEmbeddings
     {

LLama.Examples/OldVersion/Quantize.cs

@@ -4,7 +4,7 @@ using System.Linq;
 using System.Text;
 using System.Threading.Tasks;
 
-namespace LLama.Examples
+namespace LLama.Examples.Old
 {
     public class Quantize
     {

LLama.Examples/OldVersion/SaveAndLoadState.cs

@@ -5,7 +5,7 @@ using System.Text;
 using System.Threading.Tasks;
 using LLama.OldVersion;
 
-namespace LLama.Examples
+namespace LLama.Examples.Old
 {
     public class SaveAndLoadState: IDisposable
     {

LLama.Examples/Program.cs

@@ -1,7 +1,6 @@
 using LLama;
 using LLama.Common;
 using LLama.Examples;
-using LLama.Examples.Old;
 
 Console.WriteLine("======================================================================================================");
 
@@ -31,9 +30,11 @@ if(version == 1)
 
     InteractiveExecutor ex = new(new LLamaModel(new ModelParams(modelPath, contextSize: 1024, seed: 1337)));
 
+    ChatSession session = new ChatSession(ex).WithOutputTransform(new LLamaTransforms.KeywordTextOutputStreamTransform(new string[] { "User:", "Bob:" }));
+
     while (prompt != "skip")
     {
-        await foreach (var text in ex.InferAsync(prompt, new InferenceParams() { Temperature = 0.6f, AntiPrompts = new List<string> { "User:" } }, default(CancellationToken)))
+        await foreach (var text in session.ChatAsync(prompt, new InferenceParams() { Temperature = 0.6f, AntiPrompts = new List<string> { "User:" } }, default(CancellationToken)))
         {
             Console.Write(text);
         }
@@ -71,5 +72,5 @@ if(version == 1)
 }
 else
 {
-    OldTestRunner.Run();
+    LLama.Examples.Old.OldTestRunner.Run();
 }

LLama/ChatSession.cs

@@ -8,15 +8,14 @@ namespace LLama
 {
     public class ChatSession
     {
-        private readonly string defaultUserName = "User";
-        private readonly string defaultAssistantName = "Assistant";
-        private readonly string defaultSystemName = "System";
-        private readonly string defaultUnknownName = "??";
         private ILLamaExecutor _executor;
         private ChatHistory _history;
         public ILLamaExecutor Executor => _executor;
         public ChatHistory History => _history;
         public SessionParams Params { get; set; }
+        public IHistoryTransform HistoryTransform { get; set; } = new LLamaTransforms.DefaultHistoryTransform();
+        public List<ITextTransform> InputTransformPipeline { get; set; } = new();
+        public ITextStreamTransform OutputTransform = new LLamaTransforms.EmptyTextOutputStreamTransform();
 
         public ChatSession(ILLamaExecutor executor, SessionParams? sessionParams = null)
         {
@@ -25,49 +24,22 @@ namespace LLama
             Params = sessionParams ?? new SessionParams();
         }
 
-        public virtual string BuildTextFromHistory(ChatHistory history)
+        public ChatSession WithHistoryTransform(IHistoryTransform transform)
         {
-            StringBuilder sb = new();
-            var userName = Params.UserName ?? defaultUserName;
-            var assistantName = Params.AssistantName ?? defaultAssistantName;
-            var systemName = Params.SystemName ?? defaultSystemName;
-            foreach (var message in history.Messages)
-            {
-                if (message.AuthorRole == AuthorRole.User)
-                {
-                    sb.AppendLine($"{userName}: {message.Content}");
-                }
-                else if (message.AuthorRole == AuthorRole.System)
-                {
-                    sb.AppendLine($"{systemName}: {message.Content}");
-                }
-                else if (message.AuthorRole == AuthorRole.Unknown)
-                {
-                    sb.AppendLine($"{defaultUnknownName}: {message.Content}");
-                }
-                else if (message.AuthorRole == AuthorRole.Assistant)
-                {
-                    sb.AppendLine($"{assistantName}: {message.Content}");
-                }
-            }
-            return sb.ToString();
+            HistoryTransform = transform;
+            return this;
         }
 
-        public virtual string CropNameFromText(string text, AuthorRole role)
+        public ChatSession AddInputTransform(ITextTransform transform)
         {
-            if (!string.IsNullOrEmpty(Params.UserName) && role == AuthorRole.User && text.StartsWith($"{Params.UserName}:"))
-            {
-                text = text.Substring($"{Params.UserName}:".Length).TrimStart();
-            }
-            else if (!string.IsNullOrEmpty(Params.AssistantName) && role == AuthorRole.Assistant && text.EndsWith($"{Params.AssistantName}:"))
-            {
-                text = text.Substring(0, text.Length - $"{Params.AssistantName}:".Length).TrimEnd();
-            }
-            if (_executor is InstructExecutor && role == AuthorRole.Assistant && text.EndsWith("\n> "))
-            {
-                text = text.Substring(0, text.Length - "\n> ".Length).TrimEnd();
-            }
-            return text;
+            InputTransformPipeline.Add(transform);
+            return this;
+        }
+
+        public ChatSession WithOutputTransform(ITextStreamTransform transform)
+        {
+            OutputTransform = transform;
+            return this;
         }
 
         /// <summary>
@@ -78,15 +50,15 @@ namespace LLama
         /// <returns></returns>
         public IEnumerable<string> Chat(ChatHistory history, InferenceParams? inferenceParams = null, CancellationToken cancellationToken = default)
         {
-            var prompt = BuildTextFromHistory(history);
-            History.AddMessage(AuthorRole.User, prompt);
+            var prompt = HistoryTransform.HistoryToText(history);
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.User, prompt).Messages);
             StringBuilder sb = new();
-            foreach (var result in _executor.Infer(prompt, inferenceParams, cancellationToken))
+            foreach (var result in ChatInternal(prompt, inferenceParams, cancellationToken))
             {
                 yield return result;
                 sb.Append(result);
             }
-            History.AddMessage(AuthorRole.Assistant, CropNameFromText(sb.ToString(), AuthorRole.Assistant));
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.Assistant, sb.ToString()).Messages);
         }
 
         /// <summary>
@@ -98,14 +70,18 @@ namespace LLama
         /// <returns></returns>
         public IEnumerable<string> Chat(string prompt, InferenceParams? inferenceParams = null, CancellationToken cancellationToken = default)
         {
-            History.AddMessage(AuthorRole.User, prompt);
+            foreach(var inputTransform in InputTransformPipeline)
+            {
+                prompt = inputTransform.Transform(prompt);
+            }
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.User, prompt).Messages);
             StringBuilder sb = new();
-            foreach (var result in _executor.Infer(prompt, inferenceParams, cancellationToken))
+            foreach (var result in ChatInternal(prompt, inferenceParams, cancellationToken))
             {
                 yield return result;
                 sb.Append(result);
             }
-            History.AddMessage(AuthorRole.Assistant, CropNameFromText(sb.ToString(), AuthorRole.Assistant));
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.Assistant, sb.ToString()).Messages);
         }
 
         /// <summary>
@@ -116,30 +92,46 @@ namespace LLama
         /// <returns></returns>
         public async IAsyncEnumerable<string> ChatAsync(ChatHistory history, InferenceParams? inferenceParams = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
-            var prompt = BuildTextFromHistory(history);
-            History.AddMessage(AuthorRole.User, prompt);
+            var prompt = HistoryTransform.HistoryToText(history);
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.User, prompt).Messages);
             StringBuilder sb = new();
-            await foreach (var result in _executor.InferAsync(prompt, inferenceParams, cancellationToken))
+            await foreach (var result in ChatAsyncInternal(prompt, inferenceParams, cancellationToken))
             {
                 yield return result;
                 sb.Append(result);
             }
-            History.AddMessage(AuthorRole.Assistant, CropNameFromText(sb.ToString(), AuthorRole.Assistant));
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.Assistant, sb.ToString()).Messages);
         }
 
         public async IAsyncEnumerable<string> ChatAsync(string prompt, InferenceParams? inferenceParams = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
         {
-            History.AddMessage(AuthorRole.User, prompt);
+            foreach (var inputTransform in InputTransformPipeline)
+            {
+                prompt = inputTransform.Transform(prompt);
+            }
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.User, prompt).Messages);
             StringBuilder sb = new();
-            await foreach (var result in _executor.InferAsync(prompt, inferenceParams, cancellationToken))
+            await foreach (var result in ChatAsyncInternal(prompt, inferenceParams, cancellationToken))
             {
                 yield return result;
                 sb.Append(result);
             }
-            History.AddMessage(AuthorRole.Assistant, CropNameFromText(sb.ToString(), AuthorRole.Assistant));
+            History.Messages.AddRange(HistoryTransform.TextToHistory(AuthorRole.Assistant, sb.ToString()).Messages);
         }
-    }
-
 
+        private IEnumerable<string> ChatInternal(string prompt, InferenceParams? inferenceParams = null, CancellationToken cancellationToken = default)
+        {
+            var results = _executor.Infer(prompt, inferenceParams, cancellationToken);
+            return OutputTransform.Transform(results);
+        }
 
+        private async IAsyncEnumerable<string> ChatAsyncInternal(string prompt, InferenceParams? inferenceParams = null, [EnumeratorCancellation] CancellationToken cancellationToken = default)
+        {
+            var results = _executor.InferAsync(prompt, inferenceParams, cancellationToken);
+            await foreach (var item in OutputTransform.TransformAsync(results))
+            {
+                yield return item;
+            }
+        }
+    }
 }

LLama/Common/SessionParams.cs

@@ -9,5 +9,19 @@ namespace LLama.Common
         public string? UserName { get; set; }
         public string? AssistantName { get; set; }
         public string? SystemName { get; set; }
+        /// <summary>
+        /// The prefix of input text. Note that this only works when you 
+        /// use the API with text as input.
+        /// </summary>
+        public string? InputPrefix { get; set; }
+        /// <summary>
+        /// The suffix of input text. Note that this only works when you 
+        /// use the API with text as input.
+        /// </summary>
+        public string? InputSuffix { get; set; }
+        /// <summary>
+        /// Whether to trim the names from the text output at the start and end.
+        /// </summary>
+        public bool TrimNamesFromOutput { get; set; } = false;
     }
 }

LLama/IHistoryTransform.cs

@@ -0,0 +1,20 @@
+using LLama.Common;
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace LLama
+{
+    public interface IHistoryTransform
+    {
+        string HistoryToText(ChatHistory history);
+        /// <summary>
+        /// 
+        /// </summary>
+        /// <param name="history">The existing history.</param>
+        /// <param name="role"></param>
+        /// <param name="text"></param>
+        /// <returns>The updated history.</returns>
+        ChatHistory TextToHistory(AuthorRole role, string text);
+    }
+}

LLama/ITextStreamTransform.cs

@@ -0,0 +1,12 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace LLama
+{
+    public interface ITextStreamTransform
+    {
+        IEnumerable<string> Transform(IEnumerable<string> tokens);
+        IAsyncEnumerable<string> TransformAsync(IAsyncEnumerable<string> tokens);
+    }
+}

LLama/ITextTransform.cs

@@ -0,0 +1,11 @@
+using System;
+using System.Collections.Generic;
+using System.Text;
+
+namespace LLama
+{
+    public interface ITextTransform
+    {
+        string Transform(string text);
+    }
+}

LLama/LLamaModel.cs

@@ -229,7 +229,6 @@ namespace LLama
 
         public void Dispose()
         {
-            NativeApi.llama_print_timings(_ctx);
             _ctx.Dispose();
         }
     }

LLama/LLamaTransforms.cs

@@ -0,0 +1,208 @@
+using LLama.Common;
+using Microsoft.VisualBasic;
+using System;
+using System.Collections.Generic;
+using System.IO;
+using System.Linq;
+using System.Security;
+using System.Security.Cryptography.X509Certificates;
+using System.Text;
+
+namespace LLama
+{
+    public class LLamaTransforms
+    {
+        public class DefaultHistoryTransform : IHistoryTransform
+        {
+            private readonly string defaultUserName = "User";
+            private readonly string defaultAssistantName = "Assistant";
+            private readonly string defaultSystemName = "System";
+            private readonly string defaultUnknownName = "??";
+
+            string _userName;
+            string _assistantName;
+            string _systemName;
+            string _unknownName;
+            bool _isInstructMode;
+            public DefaultHistoryTransform(string? userName = null, string? assistantName = null, 
+                string? systemName = null, string? unknownName = null, bool isInstructMode = false)
+            {
+                _userName = userName ?? defaultUserName;
+                _assistantName = assistantName ?? defaultAssistantName;
+                _systemName = systemName ?? defaultSystemName;
+                _unknownName = unknownName ?? defaultUnknownName;
+                _isInstructMode = isInstructMode;
+            }
+
+            public virtual string HistoryToText(ChatHistory history)
+            {
+                StringBuilder sb = new();
+                foreach (var message in history.Messages)
+                {
+                    if (message.AuthorRole == AuthorRole.User)
+                    {
+                        sb.AppendLine($"{_userName}: {message.Content}");
+                    }
+                    else if (message.AuthorRole == AuthorRole.System)
+                    {
+                        sb.AppendLine($"{_systemName}: {message.Content}");
+                    }
+                    else if (message.AuthorRole == AuthorRole.Unknown)
+                    {
+                        sb.AppendLine($"{_unknownName}: {message.Content}");
+                    }
+                    else if (message.AuthorRole == AuthorRole.Assistant)
+                    {
+                        sb.AppendLine($"{_assistantName}: {message.Content}");
+                    }
+                }
+                return sb.ToString();
+            }
+
+            public virtual ChatHistory TextToHistory(AuthorRole role, string text)
+            {
+                ChatHistory history = new ChatHistory();
+                history.AddMessage(role, TrimNamesFromText(text, role));
+                return history;
+            }
+
+            public virtual string TrimNamesFromText(string text, AuthorRole role)
+            {
+                if (role == AuthorRole.User && text.StartsWith($"{_userName}:"))
+                {
+                    text = text.Substring($"{_userName}:".Length).TrimStart();
+                }
+                else if (role == AuthorRole.Assistant && text.EndsWith($"{_assistantName}:"))
+                {
+                    text = text.Substring(0, text.Length - $"{_assistantName}:".Length).TrimEnd();
+                }
+                if (_isInstructMode && role == AuthorRole.Assistant && text.EndsWith("\n> "))
+                {
+                    text = text.Substring(0, text.Length - "\n> ".Length).TrimEnd();
+                }
+                return text;
+            }
+        }
+
+        /// <summary>
+        /// A text input transform that only trims the text.
+        /// </summary>
+        public class NaiveTextInputTransform : ITextTransform
+        {
+            public NaiveTextInputTransform()
+            {
+                
+            }
+
+            public string Transform(string text)
+            {
+                return text.Trim();
+            }
+        }
+
+        public class EmptyTextOutputStreamTransform : ITextStreamTransform
+        {
+            public IEnumerable<string> Transform(IEnumerable<string> tokens)
+            {
+                return tokens;
+            }
+
+            public IAsyncEnumerable<string> TransformAsync(IAsyncEnumerable<string> tokens)
+            {
+                return tokens;
+            }
+        }
+
+        public class KeywordTextOutputStreamTransform : ITextStreamTransform
+        {
+            HashSet<string> _keywords;
+            int _maxKeywordLength;
+            bool _removeAllMatchedTokens;
+
+            /// <summary>
+            /// 
+            /// </summary>
+            /// <param name="keywords">Keywords that you want to remove from the response.</param>
+            /// <param name="redundancyLength">The extra length when searching for the keyword. For example, if your only keyword is "highlight", 
+            /// maybe the token you get is "\r\nhighligt". In this condition, if redundancyLength=0, the token cannot be successfully matched because the length of "\r\nhighligt" (10)
+            /// has already exceeded the maximum length of the keywords (8). On the contrary, setting redundancyLengyh >= 2 leads to successful match.
+            /// The larger the redundancyLength is, the lower the processing speed. But as an experience, it won't introduce too much performance impact when redundancyLength <= 5 </param>
+            /// <param name="removeAllMatchedTokens">If set to true, when getting a matched keyword, all the related tokens will be removed. Otherwise only the part of keyword will be removed.</param>
+            public KeywordTextOutputStreamTransform(IEnumerable<string> keywords, int redundancyLength = 3, bool removeAllMatchedTokens = false)
+            {
+                _keywords = new(keywords);
+                _maxKeywordLength = keywords.Select(x => x.Length).Max() + redundancyLength;
+                _removeAllMatchedTokens = removeAllMatchedTokens;
+            }
+
+            public IEnumerable<string> Transform(IEnumerable<string> tokens)
+            {
+                var window = new Queue<string>();
+
+                foreach (var s in tokens)
+                {
+                    window.Enqueue(s);
+                    var current = string.Join("", window);
+                    if (_keywords.Any(x => current.Contains(x)))
+                    {
+                        int total = window.Count;
+                        for (int i = 0; i < total; i++)
+                        {
+                            window.Dequeue();
+                        }
+                    }
+                    if(current.Length >= _maxKeywordLength)
+                    {
+                        int total = window.Count;
+                        for (int i = 0; i < total; i++)
+                        {
+                            yield return window.Dequeue();
+                        }
+                    }
+                }
+                int totalCount = window.Count;
+                for (int i = 0; i < totalCount; i++)
+                {
+                    yield return window.Dequeue();
+                }
+            }
+
+            public async IAsyncEnumerable<string> TransformAsync(IAsyncEnumerable<string> tokens)
+            {
+                var window = new Queue<string>();
+
+                await foreach (var s in tokens)
+                {
+                    window.Enqueue(s);
+                    var current = string.Join("", window);
+                    if (_keywords.Any(x => current.Contains(x)))
+                    {
+                        var matchedKeyword = _keywords.First(x => current.Contains(x));
+                        int total = window.Count;
+                        for (int i = 0; i < total; i++)
+                        {
+                            window.Dequeue();
+                        }
+                        if (!_removeAllMatchedTokens)
+                        {
+                            yield return current.Replace(matchedKeyword, "");
+                        }
+                    }
+                    if (current.Length >= _maxKeywordLength)
+                    {
+                        int total = window.Count;
+                        for (int i = 0; i < total; i++)
+                        {
+                            yield return window.Dequeue();
+                        }
+                    }
+                }
+                int totalCount = window.Count;
+                for (int i = 0; i < totalCount; i++)
+                {
+                    yield return window.Dequeue();
+                }
+            }
+        }
+    }
+}