中文對話式大語言模型Firefly-2b6開源，使用210萬訓(xùn)練數(shù)據(jù)

在文章Firefly(流螢): 中文對話式大語言模型中，我們介紹了關(guān)于Firefly(流螢)項目的工作，并且分享了我們訓(xùn)練的firefly-1b4模型。這是Firefly項目開源的第一個模型，雖然取得了還不錯的效果，但無論是訓(xùn)練數(shù)據(jù)還是模型參數(shù)量，都還有很大的優(yōu)化空間。

所以，在firefly-1b4實(shí)驗的基礎(chǔ)上，我們對訓(xùn)練數(shù)據(jù)進(jìn)行清洗，并且增加了數(shù)據(jù)量，得到210萬數(shù)據(jù)，并用它訓(xùn)練得到了firefly-2b6模型。

在本文中，我們將對該模型進(jìn)行分享和介紹。與firefly-1b4相比，firefly-2b6的代碼生成能力取得了較大的進(jìn)步，并且在古詩詞生成、對聯(lián)、作文、開放域生成等方面也有不錯的提升。

firefly-1b4和firefly-2b6的訓(xùn)練配置如下表所示。無論是訓(xùn)練數(shù)據(jù)量，還是訓(xùn)練步數(shù)，firefly-2b6都更加充分。

項目地址：

https://github.com/yangjianxin1/Firefly

模型權(quán)重鏈接見文末。

模型使用

使用如下代碼即可使用模型：

from transformers import BloomTokenizerFast, BloomForCausalLM
device = 'cuda'
path = 'YeungNLP/firefly-2b6'


tokenizer = BloomTokenizerFast.from_pretrained(path)
model = BloomForCausalLM.from_pretrained(path)
model.eval()
model = model.to(device)
text = input('User：')
while True:
    text = '{}'.format(text)
    input_ids = tokenizer(text, return_tensors="pt").input_ids
    input_ids = input_ids.to(device)
outputs=model.generate(input_ids,max_new_tokens=250,do_sample=True,top_p=0.7,temperature=0.35,
                             repetition_penalty=1.2, eos_token_id=tokenizer.eos_token_id)
    rets = tokenizer.batch_decode(outputs)
    output = rets[0].strip().replace(text, "").replace('', "")
    print("Firefly：{}".format(output))
    text = input('User：')

代碼生成

盡管在訓(xùn)練集中，代碼的數(shù)據(jù)量不多，但令人驚喜的是，firefly-2b6已經(jīng)具備一定的代碼生成能力。

在筆者的實(shí)測中，對于一些編程題，firefly-2b6生成的代碼可以做到無需修改，直接運(yùn)行成功，并且得到正確的答案。下面將展示一些編程題的生成例子。

示例1：幫我用python寫一個冒泡排序算法。

示例2：用python實(shí)現(xiàn)一個快速排序算法，輸入為一個數(shù)組，返回排序好之后的數(shù)組。

示例3：用python寫一個二分查找算法。

示例4：寫一個函數(shù)，計算數(shù)組中偶數(shù)的個數(shù)，輸入為數(shù)組，輸出為偶數(shù)的個數(shù)。

示例5：用html生成一個教務(wù)管理系統(tǒng)的登錄界面，要求包含用戶名、密碼輸入框和登錄按鈕。

該html代碼在瀏覽器中的效果如下圖：

其他樣例

同樣，我們也對文言文、古詩詞、文章生成等數(shù)據(jù)進(jìn)行了清洗，提高數(shù)據(jù)的質(zhì)量。實(shí)測下來，我們發(fā)現(xiàn)firefly-2b6的生成效果，確實(shí)提升了不少。

數(shù)據(jù)質(zhì)量的優(yōu)化，對文言文翻譯任務(wù)的提升，尤為明顯。在訓(xùn)練firefly-1b4時，文言文數(shù)據(jù)為較短的句子對。但在訓(xùn)練firefly-2b6時，我們使用了較長篇幅的文本對。

下面為一些實(shí)測的例子。

文章小結(jié)

雖然firefly-2b6已經(jīng)初步具備代碼生成能力，但由于訓(xùn)練集中的代碼數(shù)據(jù)的數(shù)量不多，對于一些編程題，效果不如人意。我們覺得仍有非常大的優(yōu)化空間，后續(xù)我們也將收集更多代碼數(shù)據(jù)，提升模型的代碼能力。

經(jīng)過firefly-1b4和firefly-2b6兩個模型的迭代，能明顯感受到增加數(shù)據(jù)量、提升數(shù)據(jù)質(zhì)量、增大模型參數(shù)量，對模型的提升非常大。

在前文中，我們提到，firefly-1b4在訓(xùn)練數(shù)據(jù)量、訓(xùn)練步數(shù)上都略有不足。為了探索"小"模型的效果上限，我們也將使用更多數(shù)量、更高質(zhì)量的數(shù)據(jù)對firefly-1b4進(jìn)行迭代。該項工作正在進(jìn)行。

后續(xù)，我們也將在多輪對話、增大模型參數(shù)量、模型量化等方向上進(jìn)行迭代，我們也將陸續(xù)開源訓(xùn)練代碼以及更多的訓(xùn)練數(shù)據(jù)。期待大家的意見和建議。