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:要设置为活动的适配器名称

返回值