Skip to content

[NLP] 使用 aitextgen 套件來生成文本

aitextgen 是一個封裝得很高級的 Python 套件,讓使用者只需要寫短短幾行的程式碼,便能設定好複雜的 AI 模型。它的架構採用了 OpenAI’s GPT-2 和 EleutherAI’s GPT Neo/GPT-3,也能使用預訓練模型(pre-trained model)繼續進行模型的微調。

以下我就簡單介紹,該如何使用 aitextgen 這個套件。


安裝

首先,我們需要使用 pip 工具來安裝 aitextgen

pip3 install aitextgen


如果後面執行程式碼時,有報錯顯示缺乏相關依賴dependencies),則再依報錯訊息補足。


生成文本

以下是一段生成文本的範例程式碼:

# coding: utf-8
from aitextgen import aitextgen


def main():
    ai = aitextgen()
    ai.generate(n=1, prompt="The dog", max_length=100)


if __name__ == "__main__":
    main()


Output:

The dog was given a free medical attention certificate and is now in foster care.

"I think it was a wonderful experience for me," said Deacon. "I really don't want to spend too much time thinking about my dog in the same way.

"The dogs are so wonderful. They are the perfect companions and they are going to be a great part of my life.

"We have a great relationship and I'm very happy about the fact we are going to


雖然論述有些奇怪,但產生的語法大部分都很正確。

當然,我們也有許多參數可以調整。

生成參數(Generation Parameters)

  • n: 文本生成的數量
  • max_length: 生成文本的長度(default=200;GPT-2 最長為 1024;GPT-neo 最長為 2048)
  • prompt: 啟動生成文本時的提示
  • temperature: 控制文本的瘋狂程度(default: 0.7)
  • num_beams: 如果大於 1,則執行集束搜尋(beam search)以生成更清晰的文本
  • repetition_penalty: 如果大於 1.0,則懲罰文本中的重複以避免無限循環
  • length_penalty: 如果大於 1.0,則懲罰與長度太長的文本
  • no_repeat_ngram_size: 避免給定重複的短句


生成函式(Generation Functions)

這裡我們假設 aitextgen 物件名稱為 ai

(若要使用 GPU,則可以使用 ai.to_gpu()ai = aitextgen(to_gpu=True)

  • ai.generate(): 生成並印出
  • ai.generate_one(): 生成單個文本並作為字串返回
  • ai.generate_samples(): 在指定溫度(temperatures)生成多個樣本
  • ai.generate_to_file(): 生成大量文本儲存至文件

讀取模型

預設讀取的模型是 124M 參數規模的 GPT-2。

但如果你想讀取其他模型,可以使用:

ai = aitextgen(model="EleutherAI/gpt-neo-125M")


如果想使用其他模型,可以參考 huggingface 的網頁:https://huggingface.co/EleutherAI


訓練模型

如果你想訓練一個模型,你可以參考官網的教學,下載莎士比亞戲劇文本,然後使用下方的範例程式碼:

from aitextgen.TokenDataset import TokenDataset
from aitextgen.tokenizers import train_tokenizer
from aitextgen.utils import GPT2ConfigCPU
from aitextgen import aitextgen

# The name of the downloaded Shakespeare text for training
file_name = "input.txt"

# Train a custom BPE Tokenizer on the downloaded text
# This will save one file: `aitextgen.tokenizer.json`, which contains the
# information needed to rebuild the tokenizer.
train_tokenizer(file_name)
tokenizer_file = "aitextgen.tokenizer.json"

# GPT2ConfigCPU is a mini variant of GPT-2 optimized for CPU-training
# e.g. the # of input tokens here is 64 vs. 1024 for base GPT-2.
config = GPT2ConfigCPU()

# Instantiate aitextgen using the created tokenizer and config
ai = aitextgen(tokenizer_file=tokenizer_file, config=config)

# You can build datasets for training by creating TokenDatasets,
# which automatically processes the dataset with the appropriate size.
data = TokenDataset(file_name, tokenizer_file=tokenizer_file, block_size=64)

# Train the model! It will save pytorch_model.bin periodically and after completion to the `trained_model` folder.
# On a 2020 8-core iMac, this took ~25 minutes to run.
ai.train(data, batch_size=8, num_steps=50000, generate_every=5000, save_every=5000)

# Generate text from it!
ai.generate(10, prompt="ROMEO:")

# With your trained model, you can reload the model at any time by
# providing the folder containing the pytorch_model.bin model weights + the config, and providing the tokenizer.
ai2 = aitextgen(model_folder="trained_model",
                tokenizer_file="aitextgen.tokenizer.json")

ai2.generate(10, prompt="ROMEO:")



儲存模型

儲存模型非常簡單。

ai.save()



References

想參考更多的教學與說明,可以參考官方的 GitHub 以及網站:


Read More

Leave a Reply