在构建知识库和机器学习模型时,数据集的导入和管理是至关重要的步骤。FastGPT提供了一种灵活的数据集导入功能,允许用户通过网页链接轻松导入文本数据集,从而丰富知识库的内容。
本文将深入分析FastGPT的源码,揭示如何通过网页链接导入数据集,并探讨其背后的技术实现细节。
业务操作流程
FastGPT 在知识库里面创建数据集支持多种方式,在知识库里面点击右上角的"新建/导入",选择"文本数据集"->"网页链接"进入新增页面。
第一步,填写网页链接、选择器信息。网页链接不做解释,选择器即 CSS 选择,前端同学肯定熟悉。此处支持填写多个链接,使用换行分割;CSS 选择器不支持多个。
注意:网页链接必须是静态网页,SPA页面是解析不出来的。
第二步,填写训练模式、处理方式相关信息。训练模式支持问答拆分,会将数据喂给 LLM,LLM 返回 QA 问答内容。处理方式支持自定义规则,适合预处理过的数据,QA 模式可以在这里修改 prompt。
第三步,上传。系统创建数据集并将数据拆分为 chunk 推入训练队列。
核心源码分析
创建数据集的 web 页面源码在pages/dataset/dtail/Import/diffSource/FileLink
文件中,没有什么特别的逻辑,FileLink
里面的三个组件分别对应新建的三个步骤,代码如下:
tsx
const LinkCollection = ({ activeStep, goToNext }: ImportDataComponentProps) => {
return (
<>
{activeStep === 0 && <CustomLinkImport goToNext={goToNext} />}
{activeStep === 1 && <DataProcess showPreviewChunks={false} goToNext={goToNext} />}
{activeStep === 2 && <Upload showPreviewChunks={false} />}
</>
);
};
最后一步上传文件时,调用的接口为/core/dataset/collection/create/link
,后端逻辑在pages/api/core/dataset/collection/create/link
文件中,核心代码如下:
ts
export default async function handler(req: NextApiRequest, res: NextApiResponse<any>) {
try {
await connectToDatabase();
const {
link,
trainingType = TrainingModeEnum.chunk,
chunkSize = 512,
chunkSplitter,
qaPrompt,
...body
} = req.body as LinkCreateDatasetCollectionParams;
const { teamId, tmbId, dataset } = await authDataset({
req,
authToken: true,
authApiKey: true,
datasetId: body.datasetId,
per: 'w'
});
// 1. check dataset limit
await checkDatasetLimit({
teamId,
insertLen: predictDataLimitLength(trainingType, new Array(10)),
standardPlans: getStandardSubPlan()
});
const { _id: collectionId } = await mongoSessionRun(async (session) => {
// 2. create collection
const collection = await createOneCollection({
...body,
name: link,
teamId,
tmbId,
type: DatasetCollectionTypeEnum.link,
trainingType,
chunkSize,
chunkSplitter,
qaPrompt,
rawLink: link,
session
});
// 3. create bill and start sync
const { billId } = await createTrainingBill({
teamId,
tmbId,
appName: 'core.dataset.collection.Sync Collection',
billSource: BillSourceEnum.training,
vectorModel: getVectorModel(dataset.vectorModel).name,
agentModel: getLLMModel(dataset.agentModel).name,
session
});
// load
await reloadCollectionChunks({
collection: {
...collection.toObject(),
datasetId: dataset
},
tmbId,
billId,
session
});
return collection;
});
jsonRes(res, {
data: { collectionId }
});
} catch (err) {
jsonRes(res, {
code: 500,
error: err
});
}
}
核心逻辑包含以下几个步骤:
- 鉴权:
parseHeaderCert
解析 token 获取teamId
,tmbId
等信息,authDatasetByTmbId
鉴定当前用户对知识库是否有写权限; - checkDatasetLimit: 检测是否达到团队maxDatasetSize上限,这个属于商业版的功能;这里面predictDataLimitLength 的逻辑比较奇怪,
predictDataLimitLength(trainingType, new Array(10))
定死了数组的长度,并不是根据实际数据进行预测; - 创建新的数据集:调用
createOneCollection
往 mongo 里面的datasets.collections
插入记录; - 创建账单:
createTrainingBill
应该也是商业版的功能; - reloadCollectionChunks:根据链接获取数据集进行
chunk
拆分,并创建训练任务。- 调用
urlsFetch
方法获取网页链接里面的数据并转成 md 格式,数据爬取引擎采用的是cheerio
; - 对数据按照规则进行
chunk
拆分; - 创建训练任务:问答训练模式采用
agentModel
,直接拆分使用vectorModel
; - 更新数据集的文本长度、标题等信息。
- 调用
ts
/* link collection start load data */
export const reloadCollectionChunks = async ({
collection,
tmbId,
billId,
rawText,
session
}: {
collection: CollectionWithDatasetType;
tmbId: string;
billId?: string;
rawText?: string;
session: ClientSession;
}) => {
const {
title,
rawText: newRawText,
collection: col,
isSameRawText
} = await getCollectionAndRawText({
collection,
newRawText: rawText
});
if (isSameRawText) return;
// split data
const { chunks } = splitText2Chunks({
text: newRawText,
chunkLen: col.chunkSize || 512
});
// insert to training queue
const model = await (() => {
if (col.trainingType === TrainingModeEnum.chunk) return col.datasetId.vectorModel;
if (col.trainingType === TrainingModeEnum.qa) return col.datasetId.agentModel;
return Promise.reject('Training model error');
})();
await MongoDatasetTraining.insertMany(
chunks.map((item, i) => ({
teamId: col.teamId,
tmbId,
datasetId: col.datasetId._id,
collectionId: col._id,
billId,
mode: col.trainingType,
prompt: '',
model,
q: item,
a: '',
chunkIndex: i
})),
{ session }
);
// update raw text
await MongoDatasetCollection.findByIdAndUpdate(
col._id,
{
...(title && { name: title }),
rawTextLength: newRawText.length,
hashRawText: hashStr(newRawText)
},
{ session }
);
};