Skip to content

Commit 2f14dfe

Browse files
stephentoubstephentoub
committed
Update for genai API changes + feedback
1 parent 4c2d892 commit 2f14dfe

3 files changed

Lines changed: 121 additions & 68 deletions

File tree

‎src/csharp/ChatClient.cs‎

Lines changed: 19 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,11 @@
88

99
namespace Microsoft.ML.OnnxRuntimeGenAI;
1010

11-
/// <summary>An <see cref="IChatClient"/> implementation based on ONNX Runtime GenAI.</summary>
12-
public sealed class ChatClient : IChatClient, IDisposable
11+
/// <summary>Provides an <see cref="IChatClient"/> implementation based on ONNX Runtime GenAI.</summary>
12+
public sealed partial class ChatClient : IChatClient
1313
{
14+
/// <summary>The options used to configure the instance.</summary>
15+
private readonly ChatClientConfiguration _config;
1416
/// <summary>The wrapped <see cref="Model"/>.</summary>
1517
private readonly Model _model;
1618
/// <summary>The wrapped <see cref="Tokenizer"/>.</summary>
@@ -20,8 +22,9 @@ public sealed class ChatClient : IChatClient, IDisposable
2022

2123
/// <summary>Initializes an instance of the <see cref="ChatClient"/> class.</summary>
2224
/// <param name="modelPath">The file path to the model to load.</param>
25+
/// <param name="configuration">Options used to configure the client instance.</param>
2326
/// <exception cref="ArgumentNullException"><paramref name="modelPath"/> is null.</exception>
24-
public ChatClient(string modelPath)
27+
public ChatClient(string modelPath, ChatClientConfiguration configuration)
2528
{
2629
if (modelPath is null)
2730
{
@@ -54,32 +57,12 @@ public ChatClient(Model model, bool ownsModel = true)
5457
_model = model;
5558
_tokenizer = new Tokenizer(_model);
5659

57-
Metadata = new("Microsoft.ML.OnnxRuntimeGenAI");
60+
Metadata = new("onnxruntime-genai");
5861
}
5962

6063
/// <inheritdoc/>
6164
public ChatClientMetadata Metadata { get; }
6265

63-
/// <summary>
64-
/// Gets or sets stop sequences to use during generation.
65-
/// </summary>
66-
/// <remarks>
67-
/// These will apply in addition to any stop sequences that are a part of the <see cref="ChatOptions.StopSequences"/>.
68-
/// </remarks>
69-
public IList<string> StopSequences { get; set; } =
70-
[
71-
// Default stop sequences based on Phi3
72-
"<|system|>",
73-
"<|user|>",
74-
"<|assistant|>",
75-
"<|end|>"
76-
];
77-
78-
/// <summary>
79-
/// Gets or sets a function that creates a prompt string from the chat history.
80-
/// </summary>
81-
public Func<IEnumerable<ChatMessage>, string> PromptFormatter { get; set; }
82-
8366
/// <inheritdoc/>
8467
public void Dispose()
8568
{
@@ -102,20 +85,20 @@ public async Task<ChatCompletion> CompleteAsync(IList<ChatMessage> chatMessages,
10285
StringBuilder text = new();
10386
await Task.Run(() =>
10487
{
105-
using Sequences tokens = _tokenizer.Encode(CreatePrompt(chatMessages));
88+
using Sequences tokens = _tokenizer.Encode(_config.PromptFormatter(chatMessages));
10689
using GeneratorParams generatorParams = new(_model);
10790
UpdateGeneratorParamsFromOptions(tokens[0].Length, generatorParams, options);
108-
generatorParams.SetInputSequences(tokens);
10991

11092
using Generator generator = new(_model, generatorParams);
93+
generator.AppendTokenSequences(tokens);
94+
11195
using var tokenizerStream = _tokenizer.CreateStream();
11296

11397
var completionId = Guid.NewGuid().ToString();
11498
while (!generator.IsDone())
11599
{
116100
cancellationToken.ThrowIfCancellationRequested();
117101

118-
generator.ComputeLogits();
119102
generator.GenerateNextToken();
120103

121104
ReadOnlySpan<int> outputSequence = generator.GetSequence(0);
@@ -147,20 +130,20 @@ public async IAsyncEnumerable<StreamingChatCompletionUpdate> CompleteStreamingAs
147130
throw new ArgumentNullException(nameof(chatMessages));
148131
}
149132

150-
using Sequences tokens = _tokenizer.Encode(CreatePrompt(chatMessages));
133+
using Sequences tokens = _tokenizer.Encode(_config.PromptFormatter(chatMessages));
151134
using GeneratorParams generatorParams = new(_model);
152135
UpdateGeneratorParamsFromOptions(tokens[0].Length, generatorParams, options);
153-
generatorParams.SetInputSequences(tokens);
154136

155137
using Generator generator = new(_model, generatorParams);
138+
generator.AppendTokenSequences(tokens);
139+
156140
using var tokenizerStream = _tokenizer.CreateStream();
157141

158142
var completionId = Guid.NewGuid().ToString();
159143
while (!generator.IsDone())
160144
{
161145
string next = await Task.Run(() =>
162146
{
163-
generator.ComputeLogits();
164147
generator.GenerateNextToken();
165148

166149
ReadOnlySpan<int> outputSequence = generator.GetSequence(0);
@@ -193,43 +176,7 @@ public object GetService(Type serviceType, object key = null) =>
193176
/// <summary>Gets whether the specified token is a stop sequence.</summary>
194177
private bool IsStop(string token, ChatOptions options) =>
195178
options?.StopSequences?.Contains(token) is true ||
196-
StopSequences?.Contains(token) is true;
197-
198-
/// <summary>Creates a prompt string from the supplied chat history.</summary>
199-
private string CreatePrompt(IEnumerable<ChatMessage> messages)
200-
{
201-
if (messages is null)
202-
{
203-
throw new ArgumentNullException(nameof(messages));
204-
}
205-
206-
if (PromptFormatter is not null)
207-
{
208-
return PromptFormatter(messages) ?? string.Empty;
209-
}
210-
211-
// Default formatting based on Phi3.
212-
StringBuilder prompt = new();
213-
214-
foreach (var message in messages)
215-
{
216-
foreach (var content in message.Contents)
217-
{
218-
switch (content)
219-
{
220-
case TextContent tc when !string.IsNullOrWhiteSpace(tc.Text):
221-
prompt.Append("<|").Append(message.Role.Value).Append("|>\n")
222-
.Append(tc.Text.Replace("<|end|>\n", ""))
223-
.Append("<|end|>\n");
224-
break;
225-
}
226-
}
227-
}
228-
229-
prompt.Append("<|assistant|>");
230-
231-
return prompt.ToString();
232-
}
179+
Array.IndexOf(_config.StopSequences, token) >= 0;
233180

234181
/// <summary>Updates the <paramref name="generatorParams"/> based on the supplied <paramref name="options"/>.</summary>
235182
private static void UpdateGeneratorParamsFromOptions(int numInputTokens, GeneratorParams generatorParams, ChatOptions options)
@@ -262,6 +209,11 @@ private static void UpdateGeneratorParamsFromOptions(int numInputTokens, Generat
262209
}
263210
}
264211

212+
if (options.Seed.HasValue)
213+
{
214+
generatorParams.SetSearchOption("random_seed", options.Seed.Value);
215+
}
216+
265217
if (options.AdditionalProperties is { } props)
266218
{
267219
foreach (var entry in props)

‎src/csharp/ChatClientConfiguration.cs‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
using Microsoft.Extensions.AI;
2+
using System;
3+
using System.Collections.Generic;
4+
5+
namespace Microsoft.ML.OnnxRuntimeGenAI;
6+
7+
/// <summary>Provides configuration options used when constructing a <see cref="ChatClient"/>.</summary>
8+
/// <remarks>
9+
/// Every model has different requirements for stop sequences and prompt formatting. For best results,
10+
/// the configuration should be tailored to the exact nature of the model being used. For example,
11+
/// when using a Phi3 model, a configuration like the following may be used:
12+
/// <code>
13+
/// static ChatClientConfiguration CreateForPhi3() =&gt;
14+
/// new(["&lt;|system|&gt;", "&lt;|user|&gt;", "&lt;|assistant|&gt;", "&lt;|end|&gt;"],
15+
/// (IEnumerable&lt;ChatMessage&gt; messages) =&gt;
16+
/// {
17+
/// StringBuilder prompt = new();
18+
///
19+
/// foreach (var message in messages)
20+
/// foreach (var content in message.Contents.OfType&lt;TextContent&gt;())
21+
/// prompt.Append("&lt;|").Append(message.Role.Value).Append("|&gt;\n").Append(tc.Text).Append("&lt;|end|&gt;\n");
22+
///
23+
/// return prompt.Append("&lt;|assistant|&gt;\n").ToString();
24+
/// });
25+
/// </code>
26+
/// </remarks>
27+
public sealed class ChatClientConfiguration
28+
{
29+
private string[] _stopSequences;
30+
private Func<IEnumerable<ChatMessage>, string> _promptFormatter;
31+
32+
/// <summary>Initializes a new instance of the <see cref="ChatClientConfiguration"/> class.</summary>
33+
/// <param name="stopSequences">The stop sequences used by the model.</param>
34+
/// <param name="promptFormatter">The function to use to format a list of messages for input into the model.</param>
35+
/// <exception cref="ArgumentNullException"><paramref name="stopSequences"/> is null.</exception>
36+
/// <exception cref="ArgumentNullException"><paramref name="promptFormatter"/> is null.</exception>
37+
public ChatClientConfiguration(
38+
string[] stopSequences,
39+
Func<IEnumerable<ChatMessage>, string> promptFormatter)
40+
{
41+
if (stopSequences is null)
42+
{
43+
throw new ArgumentNullException(nameof(stopSequences));
44+
}
45+
46+
if (promptFormatter is null)
47+
{
48+
throw new ArgumentNullException(nameof(promptFormatter));
49+
}
50+
51+
StopSequences = stopSequences;
52+
PromptFormatter = promptFormatter;
53+
}
54+
55+
/// <summary>
56+
/// Gets or sets stop sequences to use during generation.
57+
/// </summary>
58+
/// <remarks>
59+
/// These will apply in addition to any stop sequences that are a part of the <see cref="ChatOptions.StopSequences"/>.
60+
/// </remarks>
61+
public string[] StopSequences
62+
{
63+
get => _stopSequences;
64+
set => _stopSequences = value ?? throw new ArgumentNullException(nameof(value));
65+
}
66+
67+
/// <summary>Gets the function that creates a prompt string from the chat history.</summary>
68+
public Func<IEnumerable<ChatMessage>, string> PromptFormatter
69+
{
70+
get => _promptFormatter;
71+
set => _promptFormatter = value ?? throw new ArgumentNullException(nameof(value));
72+
}
73+
}

‎test/csharp/TestOnnxRuntimeGenAIAPI.cs‎

Lines changed: 29 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,9 +5,11 @@
55
using System.IO;
66
using System.Linq;
77
using System.Runtime.InteropServices;
8-
using System.Runtime.CompilerServices;
98
using Xunit;
109
using Xunit.Abstractions;
10+
using System.Collections.Generic;
11+
using Microsoft.Extensions.AI;
12+
using System.Text;
1113

1214
namespace Microsoft.ML.OnnxRuntimeGenAI.Tests
1315
{
@@ -349,6 +351,32 @@ public void TestTopKTopPSearch()
349351
}
350352
}
351353

354+
[IgnoreOnModelAbsenceFact(DisplayName = "TestChatClient")]
355+
public async void TestChatClient()
356+
{
357+
using var client = new ChatClient(
358+
_phi2Path,
359+
new(["<|system|>", "<|user|>", "<|assistant|>", "<|end|>"],
360+
(IEnumerable<ChatMessage> messages) =>
361+
{
362+
StringBuilder prompt = new();
363+
364+
foreach (var message in messages)
365+
foreach (var content in message.Contents.OfType<TextContent>())
366+
prompt.Append("<|").Append(message.Role.Value).Append("|>\n").Append(content.Text).Append("<|end|>\n");
367+
368+
return prompt.Append("<|assistant|>\n").ToString();
369+
}));
370+
371+
var completion = await client.CompleteAsync("What is 2 + 3?", new()
372+
{
373+
MaxOutputTokens = 20,
374+
Temperature = 0f,
375+
});
376+
377+
Assert.Contains("5", completion.ToString());
378+
}
379+
352380
[IgnoreOnModelAbsenceFact(DisplayName = "TestTokenizerBatchEncodeDecode")]
353381
public void TestTokenizerBatchEncodeDecode()
354382
{

0 commit comments

Comments
 (0)