中文对话式大语言模型Firefly-2b6开源，使用210万训练数据

 

 

在文章Firefly(流萤): 中文对话式大语言模型中，我们介绍了关于Firefly(流萤)项目的工作，并且分享了我们训练的firefly-1b4模型。这是Firefly项目开源的第一个模型，虽然取得了还不错的效果，但无论是训练数据还是模型参数量，都还有很大的优化空间。

 

所以，在firefly-1b4实验的基础上，我们对训练数据进行清洗，并且增加了数据量，得到210万数据，并用它训练得到了firefly-2b6模型。

 

在本文中，我们将对该模型进行分享和介绍。与firefly-1b4相比，firefly-2b6的代码生成能力取得了较大的进步，并且在古诗词生成、对联、作文、开放域生成等方面也有不错的提升。

 

firefly-1b4和firefly-2b6的训练配置如下表所示。无论是训练数据量，还是训练步数，firefly-2b6都更加充分。

参数	firefly-1b4	firefly-2b6
batch size	16	8
learning rate	3e-5	3e-5
warmup step	3000	3000
lr schedule	cosine	cosine
max length	512	512
training step	90k	260k
训练集规模	160万	210万

 

项目地址：

https://github.com/yangjianxin1/Firefly

 

模型权重链接见文末。

 

模型使用

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

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：')

 

 

代码生成

尽管在训练集中，代码的数据量不多，但令人惊喜的是，firefly-2b6已经具备一定的代码生成能力。

 

在笔者的实测中，对于一些编程题，firefly-2b6生成的代码可以做到无需修改，直接运行成功，并且得到正确的答案。下面将展示一些编程题的生成例子。

 

示例1：帮我用python写一个冒泡排序算法。

 

示例2：用python实现一个快速排序算法，输入为一个数组，返回排序好之后的数组。

示例3：用python写一个二分查找算法。

示例4：写一个函数，计算数组中偶数的个数，输入为数组，输出为偶数的个数。

示例5：用html生成一个教务管理系统的登录界面，要求包含用户名、密码输入框和登录按钮。

该html代码在浏览器中的效果如下图：

 

其他样例

同样，我们也对文言文、古诗词、文章生成等数据进行了清洗，提高数据的质量。实测下来，我们发现firefly-2b6的生成效果，确实提升了不少。

 

数据质量的优化，对文言文翻译任务的提升，尤为明显。在训练firefly-1b4时，文言文数据为较短的句子对。但在训练firefly-2b6时，我们使用了较长篇幅的文本对。

 

下面为一些实测的例子。

 

 

 

 

 

 

 

 

 

 

文章小结

虽然firefly-2b6已经初步具备代码生成能力，但由于训练集中的代码数据的数量不多，对于一些编程题，效果不如人意。我们觉得仍有非常大的优化空间，后续我们也将收集更多代码数据，提升模型的代码能力。

 

经过firefly-1b4和firefly-2b6两个模型的迭代，能明显感受到增加数据量、提升数据质量、增大模型参数量，对模型的提升非常大。

 

在前文中，我们提到，firefly-1b4在训练数据量、训练步数上都略有不足。为了探索"小"模型的效果上限，我们也将使用更多数量、更高质量的数据对firefly-1b4进行迭代。该项工作正在进行。

 

后续，我们也将在多轮对话、增大模型参数量、模型量化等方向上进行迭代，我们也将陆续开源训练代码以及更多的训练数据。期待大家的意见和建议。

 

 

---

审核编辑 ：李倩