| using System; |
| using System.IO; |
| using System.Text; |
| using System.Threading.Tasks; |
| using Microsoft.ML.OnnxRuntimeGenAI; |
| using UnityEngine; |
|
|
| namespace Sky.OnnxRuntime.Samples |
| { |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| public sealed class GenAITextGenerationSample : MonoBehaviour |
| { |
| [Tooltip("Absolute path to a folder containing an onnxruntime-genai model " + |
| "(genai_config.json, tokenizer files and the .onnx weights).")] |
| [SerializeField] |
| private string _modelFolderPath = ""; |
|
|
| [TextArea] |
| [SerializeField] |
| private string _prompt = "Write a haiku about the Unity game engine."; |
|
|
| [Tooltip("Maximum number of tokens in the prompt + response.")] |
| [SerializeField] |
| private int _maxLength = 256; |
|
|
| private void Start() |
| { |
| if (string.IsNullOrWhiteSpace(_modelFolderPath)) |
| { |
| Debug.LogWarning( |
| "[ORT GenAI] Set 'Model Folder Path' on the GenAITextGenerationSample " + |
| "component to a local onnxruntime-genai model folder, then enter Play Mode."); |
| return; |
| } |
|
|
| if (!Directory.Exists(_modelFolderPath)) |
| { |
| Debug.LogError($"[ORT GenAI] Model folder not found: {_modelFolderPath}"); |
| return; |
| } |
|
|
| string modelPath = _modelFolderPath; |
| string prompt = _prompt; |
| int maxLength = _maxLength; |
|
|
| |
| Task.Run(() => |
| { |
| try |
| { |
| Generate(modelPath, prompt, maxLength); |
| } |
| catch (Exception e) |
| { |
| Debug.LogError($"[ORT GenAI] Generation failed: {e}"); |
| } |
| }); |
| } |
|
|
| private static void Generate(string modelPath, string prompt, int maxLength) |
| { |
| Debug.Log($"[ORT GenAI] Loading model from: {modelPath}"); |
|
|
| using var model = new Model(modelPath); |
| using var tokenizer = new Tokenizer(model); |
|
|
| |
| |
| using var sequences = tokenizer.Encode(prompt); |
|
|
| using var generatorParams = new GeneratorParams(model); |
| generatorParams.SetSearchOption("max_length", maxLength); |
|
|
| using var generator = new Generator(model, generatorParams); |
| generator.AppendTokenSequences(sequences); |
|
|
| |
| using var tokenizerStream = tokenizer.CreateStream(); |
| var response = new StringBuilder(); |
|
|
| while (!generator.IsDone()) |
| { |
| generator.GenerateNextToken(); |
|
|
| ReadOnlySpan<int> newTokens = generator.GetNextTokens(); |
| if (newTokens.Length == 0) |
| continue; |
|
|
| string piece = tokenizerStream.Decode(newTokens[newTokens.Length - 1]); |
| response.Append(piece); |
| } |
|
|
| Debug.Log($"[ORT GenAI] Prompt:\n{prompt}"); |
| Debug.Log($"[ORT GenAI] Response:\n{response}"); |
| Debug.Log("[ORT GenAI] Generation complete ✅"); |
| } |
| } |
| } |
|
|