@@ -33,7 +33,7 @@ public sealed class OnnxTextGenerator : ITextGenerator, IDisposable
3333 /// Tokenizer used with the Onnx Generator and Model classes to produce tokens.
3434 /// This has the potential to contain a null value, depending on the contents of the Model Directory.
3535 /// </summary>
36- private readonly Tokenizer ? _tokenizer = default ;
36+ private readonly Tokenizer _tokenizer ;
3737
3838 /// <summary>
3939 /// Tokenizer used for GetTokens() and CountTokens()
@@ -84,15 +84,55 @@ public OnnxTextGenerator(
8484 this . _log . LogDebug ( "Onnx model loaded" ) ;
8585 }
8686
87+ /// <inheritdoc/>
88+ public int CountTokens ( string text )
89+ {
90+ // TODO: Implement with _tokenizer and remove _textTokenizer
91+ return this . _textTokenizer . CountTokens ( text ) ;
92+ }
93+
94+ /// <inheritdoc/>
95+ public IReadOnlyList < string > GetTokens ( string text )
96+ {
97+ // TODO: Implement with _tokenizer and remove _textTokenizer
98+ return this . _textTokenizer . GetTokens ( text ) ;
99+ }
100+
87101 /// <inheritdoc/>
88102 public async IAsyncEnumerable < GeneratedTextContent > GenerateTextAsync (
89103 string prompt ,
90104 TextGenerationOptions ? options = null ,
91105 [ EnumeratorCancellation ] CancellationToken cancellationToken = default )
92106 {
93- var tokens = this . _tokenizer ? . Encode ( prompt ) ;
107+ // TODO: the prompt format should be configurable
108+ using var sequences = this . _tokenizer . Encode ( $ "<|user|>{ prompt } <|end|><|assistant|>") ;
109+
94110 using var generatorParams = new GeneratorParams ( this . _model ) ;
111+ this . SetGeneratorParams ( generatorParams , options ) ;
95112
113+ using var tokenizerStream = this . _tokenizer . CreateStream ( ) ;
114+ using var generator = new Generator ( this . _model , generatorParams ) ;
115+ generator . AppendTokenSequences ( sequences ) ;
116+
117+ while ( ! generator . IsDone ( ) )
118+ {
119+ generator . GenerateNextToken ( ) ;
120+ var x = tokenizerStream . Decode ( generator . GetSequence ( 0 ) [ ^ 1 ] ) ;
121+ yield return new GeneratedTextContent ( x ) ;
122+ }
123+
124+ await Task . CompletedTask . ConfigureAwait ( false ) ;
125+ }
126+
127+ /// <inheritdoc/>
128+ public void Dispose ( )
129+ {
130+ this . _model . Dispose ( ) ;
131+ this . _tokenizer . Dispose ( ) ;
132+ }
133+
134+ private void SetGeneratorParams ( GeneratorParams generatorParams , TextGenerationOptions ? options )
135+ {
96136 generatorParams . SetSearchOption ( "max_length" , this . MaxTokenTotal ) ;
97137 generatorParams . SetSearchOption ( "min_length" , this . _config . MinLength ) ;
98138 generatorParams . SetSearchOption ( "num_return_sequences" , this . _config . ResultsPerPrompt ) ;
@@ -145,49 +185,5 @@ public async IAsyncEnumerable<GeneratedTextContent> GenerateTextAsync(
145185
146186 break ;
147187 }
148-
149- generatorParams . SetInputSequences ( tokens ) ;
150-
151- using ( var generator = new Generator ( this . _model , generatorParams ) )
152- {
153- List < int > outputTokens = [ ] ;
154-
155- while ( ! generator . IsDone ( ) && cancellationToken . IsCancellationRequested == false )
156- {
157- generator . ComputeLogits ( ) ;
158- generator . GenerateNextToken ( ) ;
159-
160- outputTokens . AddRange ( generator . GetSequence ( 0 ) ) ;
161-
162- if ( outputTokens . Count > 0 && this . _tokenizer != null )
163- {
164- var newToken = outputTokens [ ^ 1 ] ;
165- yield return this . _tokenizer . Decode ( [ newToken ] ) ;
166- }
167- }
168- }
169-
170- await Task . CompletedTask . ConfigureAwait ( false ) ;
171- }
172-
173- /// <inheritdoc/>
174- public int CountTokens ( string text )
175- {
176- // TODO: Implement with _tokenizer and remove _textTokenizer
177- return this . _textTokenizer . CountTokens ( text ) ;
178- }
179-
180- /// <inheritdoc/>
181- public IReadOnlyList < string > GetTokens ( string text )
182- {
183- // TODO: Implement with _tokenizer and remove _textTokenizer
184- return this . _textTokenizer . GetTokens ( text ) ;
185- }
186-
187- /// <inheritdoc/>
188- public void Dispose ( )
189- {
190- this . _model ? . Dispose ( ) ;
191- this . _tokenizer ? . Dispose ( ) ;
192188 }
193189}
0 commit comments