Python API
注意:此 API 处于预览阶段,可能会有所更改。
安装和导入
Python API 由 onnxruntime-genai Python 包提供。
pip install onnxruntime-genai
import onnxruntime_genai
Model 类
加载模型
从磁盘上的文件夹加载 ONNX 模型和配置。
onnxruntime_genai.Model(model_folder: str) -> onnxruntime_genai.Model
参数
model_folder
:模型和配置在磁盘上的位置
返回值
onnxruntime_genai.Model
Generate 方法
onnxruntime_genai.Model.generate(params: GeneratorParams) -> numpy.ndarray[int, int]
参数
params
:(必填)由GeneratorParams
方法创建。
返回值
numpy.ndarray[int, int]
:一个二维 numpy 数组,其维度等于传入的批次大小和令牌序列的最大长度。
设备类型
返回模型配置运行的设备类型。
onnxruntime_genai.Model.device_type
返回值
str
:一个描述加载模型将运行的设备的字符串
Tokenizer 类
创建分词器对象
onnxruntime_genai.Model.Tokenizer(model: onnxruntime_genai.Model) -> onnxruntime_genai.Tokenizer
参数
model
:(必填)通过Model()
加载的模型
返回值
Tokenizer
:分词器对象
编码
onnxruntime_genai.Tokenizer.encode(text: str) -> numpy.ndarray[numpy.int32]
参数
text
:(必填)
返回值
numpy.ndarray[numpy.int32]
:一个表示提示的令牌数组
解码
onnxruntime_genai.Tokenizer.decode(tokens: numpy.ndarry[int]) -> str
参数
numpy.ndarray[numpy.int32]
:(必填)一个生成的令牌序列
返回值
str
:已解码的生成的令牌
批量编码
onnxruntime_genai.Tokenizer.encode_batch(texts: list[str]) -> numpy.ndarray[int, int]
参数
texts
:输入列表
返回值
numpy.ndarray[int, int]
:已分词字符串的批次
批量解码
onnxruntime_genai.Tokenize.decode_batch(tokens: [[numpy.int32]]) -> list[str]
参数
- tokens
返回值
texts
:一批解码文本
创建分词器解码流
onnxruntime_genai.Tokenizer.create_stream() -> TokenizerStream
参数
无
返回值
onnxruntime_genai.TokenizerStream
分词器流对象
TokenizerStream 类
此类的作用是累积下一个可显示字符串(根据分词器的词汇表)。
解码方法
onnxruntime_genai.TokenizerStream.decode(token: int32) -> str
参数
token
:(必填)要解码的令牌
返回值
str
:如果已累积可显示字符串,此方法返回该字符串。否则,此方法返回空字符串。
GeneratorParams 类
创建一个 Generator 参数对象
onnxruntime_genai.GeneratorParams(model: Model) -> GeneratorParams
填充令牌 ID 成员
onnxruntime_genai.GeneratorParams.pad_token_id
EOS 令牌 ID 成员
onnxruntime_genai.GeneratorParams.eos_token_id
词汇量大小成员
onnxruntime_genai.GeneratorParams.vocab_size
input_ids 成员
onnxruntime_genai.GeneratorParams.input_ids: numpy.ndarray[numpy.int32, numpy.int32]
设置模型输入
onnxruntime_genai.GeneratorParams.set_model_input(name: str, value: [])
设置搜索选项方法
onnxruntime_genai.GeneratorParams.set_search_options(options: dict[str, Any])
尝试使用最大批处理大小进行图捕获
onnxruntime_genai.GeneratorParams.try_graph_capture_with_max_batch_size(max_batch_size: int)
Generator 类
创建一个 Generator
onnxruntime_genai.Generator(model: Model, params: GeneratorParams) -> Generator
参数
model
:(必填)用于生成的模型params
:(必填)控制生成的一组参数
返回值
onnxruntime_genai.Generator
Generator 对象
生成是否完成
onnxruntime_genai.Generator.is_done() -> bool
返回值
当所有序列达到最大长度或已达到序列末尾时返回 true。
计算 logits
模型运行一次迭代。
onnxruntime_genai.Generator.compute_logits()
获取输出
返回模型的输出。
onnxruntime_genai.Generator.get_output(str: name) -> numpy.ndarray
参数
name
:模型输出的名称
返回值
numpy.ndarray
:一个多维数组,表示模型输出。数组的形状是输出的形状。
示例
以下代码返回模型的输出 logits。
logits = generator.get_output("logits")
生成下一个令牌
使用当前 logits 集和指定的生成器参数,通过 Top P 采样计算下一批令牌。
onnxruntime_genai.Generator.generate_next_token()
获取下一个令牌
onnxruntime_genai.Generator.get_next_tokens() -> numpy.ndarray[numpy.int32]
返回值
numpy.ndarray[numpy.int32]
:最新生成的令牌
获取序列
onnxruntime_genai.Generator.get_sequence(index: int) -> numpy.ndarray[numpy.int32]
index
:(必填)要返回的批次中序列的索引
Adapter 类
创建
使用已加载的模型创建一个 Adapters 对象。
model = ...
adapters = og.Adapters(model)
参数
model
:将与适配器一起使用的模型
返回值
一个 Adapter
对象
加载
将适配器从磁盘加载到内存中的 Adapter 对象。
onnxruntime_genai.Adapters(file: str, name: str) -> None
参数
file
:从磁盘加载适配器的位置name
:适配器的名称
返回值
无
设置活动适配器
在 Generator
对象上设置活动适配器。
onnxruntime_genai.Generator(adapters: Generators::Adapters, adapter: str) -> None
参数
adapters
:适配器对象,已将标识的适配器加载到其中adapter
:要设置为活动的适配器名称
返回值
无