Skip to content

使用 Outlines 工具結構化模型的輸出

Last Updated on 2024-09-02 by Clay

在將大型語言模型(Large Language Model, LLM)應用於實際場景時,經常不僅僅是讓模型自由發揮、任意生成文句 —— 我們也可能會希望模型返回特定的結構,比方說做選擇題、或是給一個評分。在這樣的情況下,transformers 架構的模型可以直接使用 outlines 這個工具。

最早我是透過同事介紹 vLLM 內部的機制時才知道 Outlines 這項工具的:Outlines 允許我們設定多個候選選項,並從中選擇模型的輸出;除此之外,我們甚至可能指定生成的是整數、浮點數、JSON 格式...... 甚至是要求要符合某個正規表示法(Regular Expression, RE)的文本生成!

能夠精準地控制 LLM 生成的格式,才好進一步將 LLM 佈署於真正應用的場景。


Outlines 的原理

Outlines 利用正規表示法定義了有限狀態機,用來限制模型生成的每一個 Token,也即是說,不符合定義好規則的 Token 不會真正被解碼。我們可以想像一個簡易的情境:

現在我們問模型:這個冰淇淋吃起來怎麼樣?然後只設定兩個候選 Tokens 讓模型選擇:Hot / Cold。

雖然即便模型正常解碼我們會得到:I'm sorry, I am a robot, I can't eat any ice cream! 但是在實際解碼的過程中,即便 I'm 這個 Token 的機率分數是 0.9,而 Hot 跟 Cold 的機率分別是 0.01 和 0.02,但由於只有 Hot 和 Cold 被允許解碼,所以模型只能產生 Cold 這個 Token。


實際應用

首先,我們可以直接透過 Python 安裝 outlines。

pip3 install outlines


接著我們可以實際讓模型做做看選擇題:

import outlines

model = outlines.models.transformers("./models/google--gemma-2-2b-it")

prompt = """How does this ice cream taste?"""

generator = outlines.generate.choice(model, ["Hot", "Cold"])
answer = generator(prompt)
print(answer)


Output:

Cold


然後就像剛剛提到的,我們可以自由地操控模型輸出的數字類別:

prompt = "<s>result of 9 + 9 = 18</s><s>result of 1 + 2 = "
answer = outlines.generate.format(model, int)(prompt)
print(answer)


Output:

3


但是如果我們把 int 換成 float

prompt = "<s>result of 9 + 9 = 18</s><s>result of 1 + 2 = "
answer = outlines.generate.format(model, float)(prompt)
print(answer)


Output:

3.0

馬上就從 3 變成了 3.0!


最後我們再來玩一個好玩的:人物設定

我們可以透過 pydantic 加上一些限制與約束,比方說名字長度、頭髮顏色選項... 等等。

from enum import Enum
from pydantic import BaseModel, constr

import outlines


class Colors(str, Enum):
    red = "red"
    brown = "brown"
    white = "white"
    black = "black"


class Character(BaseModel):
    name: constr(max_length=10)
    age: int
    hair_color: Colors


generator = outlines.generate.json(model, Character)
character = generator("Give me a character description")
print(character)


Output:

name='The Oracle' age=32 hair_color=<Colors.black: 'black'>

話說居然是生成 The Oracle 嗎... 難道是駭客任務(The Matrix)看太多了嗎?不過 32 歲也太年輕了吧。


當然我們可以再修改一下 Prompt,讓其生成我們想要的角色:

generator = outlines.generate.json(model, Character)
character = generator("Give me a superman character description")
print(character)


Output:

總之,這個工具可以應用在各個方面,它也支援像是 OLLAMA、vLLM、llama-cpp 等知名框架,算是整合性非常好的工具了,總覺得它還有許多潛力可以應用在各種服務上呢!


References


Read More

Leave a Reply