488 lines
16 KiB
JavaScript
488 lines
16 KiB
JavaScript
'use client';
|
||
|
||
import { useTranslation } from 'react-i18next';
|
||
import { toast } from 'sonner';
|
||
import axios from 'axios';
|
||
|
||
const useDatasetExport = projectId => {
|
||
const { t } = useTranslation();
|
||
|
||
// 优化的流式导出 - 使用 WritableStream 避免内存溢出
|
||
const exportDatasetsStreaming = async (exportOptions, onProgress) => {
|
||
try {
|
||
const batchSize = exportOptions.batchSize || 1000;
|
||
let offset = 0;
|
||
let hasMore = true;
|
||
let totalProcessed = 0;
|
||
let isFirstBatch = true;
|
||
|
||
// 确定文件格式
|
||
const fileFormat = exportOptions.fileFormat || 'json';
|
||
const formatType = exportOptions.formatType || 'alpaca';
|
||
|
||
// 生成文件名
|
||
const formatSuffixMap = {
|
||
alpaca: 'alpaca',
|
||
multilingualthinking: 'multilingual-thinking',
|
||
sharegpt: 'sharegpt',
|
||
custom: 'custom'
|
||
};
|
||
const formatSuffix = formatSuffixMap[formatType] || formatType || 'export';
|
||
const balanceSuffix = exportOptions.balanceMode ? '-balanced' : '';
|
||
const dateStr = new Date().toISOString().slice(0, 10);
|
||
const fileName = `datasets-${projectId}-${formatSuffix}${balanceSuffix}-${dateStr}.${fileFormat}`;
|
||
|
||
// 创建可写流
|
||
let fileStream;
|
||
let writer;
|
||
|
||
try {
|
||
// 使用 showSaveFilePicker API（现代浏览器）
|
||
if (window.showSaveFilePicker) {
|
||
const handle = await window.showSaveFilePicker({
|
||
suggestedName: fileName,
|
||
types: [
|
||
{
|
||
description: 'Dataset File',
|
||
accept: {
|
||
'application/json': [`.${fileFormat}`]
|
||
}
|
||
}
|
||
]
|
||
});
|
||
fileStream = await handle.createWritable();
|
||
} else {
|
||
// 降级方案：使用内存缓冲区（但分块处理）
|
||
fileStream = null;
|
||
}
|
||
} catch (err) {
|
||
// 用户取消或不支持，使用降级方案
|
||
fileStream = null;
|
||
}
|
||
|
||
// 如果不支持流式写入，使用分块累积方案
|
||
let chunks = [];
|
||
let chunkCount = 0;
|
||
const MAX_CHUNKS_IN_MEMORY = 5; // 最多在内存中保留5批数据
|
||
|
||
// 写入文件头（JSON数组开始或CSV表头）
|
||
if (fileFormat === 'json') {
|
||
if (fileStream) {
|
||
await fileStream.write('[\n');
|
||
} else {
|
||
chunks.push('[\n');
|
||
}
|
||
} else if (fileFormat === 'csv') {
|
||
// 写入CSV表头
|
||
const headers = getCSVHeaders(formatType, exportOptions);
|
||
const headerLine = headers.join(',') + '\n';
|
||
if (fileStream) {
|
||
await fileStream.write(headerLine);
|
||
} else {
|
||
chunks.push(headerLine);
|
||
}
|
||
}
|
||
|
||
// 分批获取和写入数据
|
||
while (hasMore) {
|
||
const apiUrl = `/api/projects/${projectId}/datasets/export`;
|
||
const requestBody = {
|
||
batchMode: true,
|
||
offset: offset,
|
||
batchSize: batchSize
|
||
};
|
||
|
||
// 如果有选中的数据集 ID，传递 ID 列表
|
||
if (exportOptions.selectedIds && exportOptions.selectedIds.length > 0) {
|
||
requestBody.selectedIds = exportOptions.selectedIds;
|
||
} else if (exportOptions.confirmedOnly) {
|
||
requestBody.status = 'confirmed';
|
||
}
|
||
|
||
// 检查是否是平衡导出模式
|
||
if (exportOptions.balanceMode && exportOptions.balanceConfig) {
|
||
requestBody.balanceMode = true;
|
||
requestBody.balanceConfig = exportOptions.balanceConfig;
|
||
}
|
||
|
||
const response = await axios.post(apiUrl, requestBody);
|
||
const batchResult = response.data;
|
||
|
||
// 如果需要包含文本块内容，批量查询并填充
|
||
if (exportOptions.customFields?.includeChunk && batchResult.data.length > 0) {
|
||
const chunkNames = batchResult.data.map(item => item.chunkName).filter(name => name);
|
||
|
||
if (chunkNames.length > 0) {
|
||
try {
|
||
const chunkResponse = await axios.post(`/api/projects/${projectId}/chunks/batch-content`, {
|
||
chunkNames
|
||
});
|
||
const chunkContentMap = chunkResponse.data;
|
||
|
||
batchResult.data.forEach(item => {
|
||
if (item.chunkName && chunkContentMap[item.chunkName]) {
|
||
item.chunkContent = chunkContentMap[item.chunkName];
|
||
}
|
||
});
|
||
} catch (chunkError) {
|
||
console.error('获取文本块内容失败:', chunkError);
|
||
}
|
||
}
|
||
}
|
||
|
||
// 转换当前批次数据
|
||
const formattedBatch = formatDataBatch(batchResult.data, exportOptions);
|
||
|
||
// 写入当前批次
|
||
if (fileFormat === 'json') {
|
||
// 保持与原逻辑一致：JSON 导出为“格式化后的 JSON 数组”（2空格缩进）
|
||
// 每条记录单独 stringify + 缩进，并在数组级别拼接，避免一次性 stringify 全量数据导致内存暴涨
|
||
const batchContent = formattedBatch
|
||
.map(item => {
|
||
const pretty = JSON.stringify(item, null, 2);
|
||
// 将对象的每一行整体再缩进 2 个空格，以符合数组元素缩进
|
||
return ' ' + pretty.replace(/\n/g, '\n ');
|
||
})
|
||
.join(',\n');
|
||
|
||
const content = isFirstBatch ? batchContent : ',\n' + batchContent;
|
||
|
||
if (fileStream) {
|
||
await fileStream.write(content);
|
||
} else {
|
||
chunks.push(content);
|
||
chunkCount++;
|
||
}
|
||
} else if (fileFormat === 'jsonl') {
|
||
const batchContent = formattedBatch.map(item => JSON.stringify(item)).join('\n') + '\n';
|
||
|
||
if (fileStream) {
|
||
await fileStream.write(batchContent);
|
||
} else {
|
||
chunks.push(batchContent);
|
||
chunkCount++;
|
||
}
|
||
} else if (fileFormat === 'csv') {
|
||
const batchContent = formatBatchToCSV(formattedBatch, formatType, exportOptions);
|
||
|
||
if (fileStream) {
|
||
await fileStream.write(batchContent);
|
||
} else {
|
||
chunks.push(batchContent);
|
||
chunkCount++;
|
||
}
|
||
}
|
||
|
||
// 如果使用内存缓冲且累积了足够多的块，触发部分下载
|
||
if (!fileStream && chunkCount >= MAX_CHUNKS_IN_MEMORY) {
|
||
// 这里我们仍然需要等到最后才能下载，但至少限制了内存使用
|
||
// 可以考虑使用 Blob 分片
|
||
}
|
||
|
||
hasMore = batchResult.hasMore;
|
||
offset = batchResult.offset;
|
||
totalProcessed += batchResult.data.length;
|
||
isFirstBatch = false;
|
||
|
||
// 通知进度更新
|
||
if (onProgress) {
|
||
onProgress({
|
||
processed: totalProcessed,
|
||
currentBatch: batchResult.data.length,
|
||
hasMore
|
||
});
|
||
}
|
||
|
||
// 避免过快请求
|
||
if (hasMore) {
|
||
await new Promise(resolve => setTimeout(resolve, 50));
|
||
}
|
||
}
|
||
|
||
// 写入文件尾
|
||
if (fileFormat === 'json') {
|
||
if (fileStream) {
|
||
await fileStream.write('\n]\n');
|
||
await fileStream.close();
|
||
} else {
|
||
chunks.push('\n]\n');
|
||
}
|
||
} else {
|
||
if (fileStream) {
|
||
await fileStream.close();
|
||
}
|
||
}
|
||
|
||
// 如果使用内存缓冲方案，现在触发下载
|
||
if (!fileStream) {
|
||
downloadFromChunks(chunks, fileName);
|
||
}
|
||
|
||
toast.success(t('datasets.exportSuccess'));
|
||
return true;
|
||
} catch (error) {
|
||
console.error('Streaming export failed:', error);
|
||
toast.error(error.message || t('datasets.exportFailed'));
|
||
return false;
|
||
}
|
||
};
|
||
|
||
// 从内存块下载文件（优化版本，使用 Blob 流）
|
||
const downloadFromChunks = (chunks, fileName) => {
|
||
// 使用 Blob 构造函数，它会自动处理大数据
|
||
const blob = new Blob(chunks, { type: 'application/octet-stream' });
|
||
const url = URL.createObjectURL(blob);
|
||
const a = document.createElement('a');
|
||
a.href = url;
|
||
a.download = fileName;
|
||
document.body.appendChild(a);
|
||
a.click();
|
||
document.body.removeChild(a);
|
||
|
||
// 延迟释放 URL，确保下载开始
|
||
setTimeout(() => URL.revokeObjectURL(url), 1000);
|
||
};
|
||
|
||
// 获取CSV表头
|
||
const getCSVHeaders = (formatType, exportOptions) => {
|
||
if (formatType === 'alpaca') {
|
||
return ['instruction', 'input', 'output', 'system'];
|
||
} else if (formatType === 'sharegpt') {
|
||
return ['messages'];
|
||
} else if (formatType === 'multilingualthinking') {
|
||
return ['reasoning_language', 'developer', 'user', 'analysis', 'final', 'messages'];
|
||
} else if (formatType === 'custom') {
|
||
const { questionField, answerField, cotField, includeLabels, includeChunk, questionOnly } =
|
||
exportOptions.customFields;
|
||
const headers = [questionField];
|
||
if (!questionOnly) {
|
||
headers.push(answerField);
|
||
if (exportOptions.includeCOT && cotField) {
|
||
headers.push(cotField);
|
||
}
|
||
}
|
||
if (includeLabels) headers.push('label');
|
||
if (includeChunk) headers.push('chunk');
|
||
return headers;
|
||
}
|
||
return [];
|
||
};
|
||
|
||
// 格式化数据批次
|
||
const formatDataBatch = (dataBatch, exportOptions) => {
|
||
const formatType = exportOptions.formatType || 'alpaca';
|
||
|
||
if (formatType === 'alpaca') {
|
||
if (exportOptions.alpacaFieldType === 'instruction') {
|
||
return dataBatch.map(({ question, answer, cot }) => ({
|
||
instruction: question,
|
||
input: '',
|
||
output: cot && exportOptions.includeCOT ? `<think>${cot}</think>\n${answer}` : answer,
|
||
system: exportOptions.systemPrompt || ''
|
||
}));
|
||
} else {
|
||
return dataBatch.map(({ question, answer, cot }) => ({
|
||
instruction: exportOptions.customInstruction || '',
|
||
input: question,
|
||
output: cot && exportOptions.includeCOT ? `<think>${cot}</think>\n${answer}` : answer,
|
||
system: exportOptions.systemPrompt || ''
|
||
}));
|
||
}
|
||
} else if (formatType === 'sharegpt') {
|
||
return dataBatch.map(({ question, answer, cot }) => {
|
||
const messages = [];
|
||
if (exportOptions.systemPrompt) {
|
||
messages.push({ role: 'system', content: exportOptions.systemPrompt });
|
||
}
|
||
messages.push({
|
||
role: 'user',
|
||
content: question
|
||
});
|
||
messages.push({
|
||
role: 'assistant',
|
||
content: cot && exportOptions.includeCOT ? `<think>${cot}</think>\n${answer}` : answer
|
||
});
|
||
return { messages };
|
||
});
|
||
} else if (formatType === 'multilingualthinking') {
|
||
return dataBatch.map(({ question, answer, cot }) => ({
|
||
reasoning_language: exportOptions.reasoningLanguage || 'English',
|
||
developer: exportOptions.systemPrompt || '',
|
||
user: question,
|
||
analysis: exportOptions.includeCOT && cot ? cot : null,
|
||
final: answer,
|
||
messages: [
|
||
{
|
||
content: exportOptions.systemPrompt || '',
|
||
role: 'system',
|
||
thinking: null
|
||
},
|
||
{
|
||
content: question,
|
||
role: 'user',
|
||
thinking: null
|
||
},
|
||
{
|
||
content: answer,
|
||
role: 'assistant',
|
||
thinking: exportOptions.includeCOT && cot ? cot : null
|
||
}
|
||
]
|
||
}));
|
||
} else if (formatType === 'custom') {
|
||
const { questionField, answerField, cotField, includeLabels, includeChunk, questionOnly } =
|
||
exportOptions.customFields;
|
||
return dataBatch.map(({ question, answer, cot, questionLabel: labels, chunkContent }) => {
|
||
const item = { [questionField]: question };
|
||
if (!questionOnly) {
|
||
item[answerField] = answer;
|
||
if (cot && exportOptions.includeCOT && cotField) {
|
||
item[cotField] = cot;
|
||
}
|
||
}
|
||
if (includeLabels && labels && labels.length > 0) {
|
||
item.label = labels.split(' ')[1];
|
||
}
|
||
if (includeChunk && chunkContent) {
|
||
item.chunk = chunkContent;
|
||
}
|
||
return item;
|
||
});
|
||
}
|
||
return dataBatch;
|
||
};
|
||
|
||
// 将批次格式化为CSV行
|
||
const formatBatchToCSV = (formattedBatch, formatType, exportOptions) => {
|
||
const headers = getCSVHeaders(formatType, exportOptions);
|
||
return (
|
||
formattedBatch
|
||
.map(item => {
|
||
return headers
|
||
.map(header => {
|
||
let field = item[header]?.toString() || '';
|
||
// 对于复杂对象，转换为JSON字符串
|
||
if (typeof item[header] === 'object') {
|
||
field = JSON.stringify(item[header]);
|
||
}
|
||
// CSV转义
|
||
if (field.includes(',') || field.includes('\n') || field.includes('"')) {
|
||
field = `"${field.replace(/"/g, '""')}"`;
|
||
}
|
||
return field;
|
||
})
|
||
.join(',');
|
||
})
|
||
.join('\n') + '\n'
|
||
);
|
||
};
|
||
|
||
// 处理和下载数据的通用函数（保留用于小数据量）
|
||
const processAndDownloadData = async (dataToExport, exportOptions) => {
|
||
const formattedData = formatDataBatch(dataToExport, exportOptions);
|
||
|
||
let content;
|
||
let fileExtension;
|
||
const fileFormat = exportOptions.fileFormat || 'json';
|
||
|
||
if (fileFormat === 'jsonl') {
|
||
content = formattedData.map(item => JSON.stringify(item)).join('\n');
|
||
fileExtension = 'jsonl';
|
||
} else if (fileFormat === 'csv') {
|
||
const headers = getCSVHeaders(exportOptions.formatType, exportOptions);
|
||
const csvRows = [
|
||
headers.join(','),
|
||
...formattedData.map(item =>
|
||
headers
|
||
.map(header => {
|
||
let field = item[header]?.toString() || '';
|
||
if (typeof item[header] === 'object') {
|
||
field = JSON.stringify(item[header]);
|
||
}
|
||
if (field.includes(',') || field.includes('\n') || field.includes('"')) {
|
||
field = `"${field.replace(/"/g, '""')}"`;
|
||
}
|
||
return field;
|
||
})
|
||
.join(',')
|
||
)
|
||
];
|
||
content = csvRows.join('\n');
|
||
fileExtension = 'csv';
|
||
} else {
|
||
content = JSON.stringify(formattedData, null, 2);
|
||
fileExtension = 'json';
|
||
}
|
||
|
||
const blob = new Blob([content], { type: 'application/json' });
|
||
const url = URL.createObjectURL(blob);
|
||
const a = document.createElement('a');
|
||
a.href = url;
|
||
|
||
const formatSuffixMap = {
|
||
alpaca: 'alpaca',
|
||
multilingualthinking: 'multilingual-thinking',
|
||
sharegpt: 'sharegpt',
|
||
custom: 'custom'
|
||
};
|
||
const formatSuffix = formatSuffixMap[exportOptions.formatType] || exportOptions.formatType || 'export';
|
||
const balanceSuffix = exportOptions.balanceMode ? '-balanced' : '';
|
||
const dateStr = new Date().toISOString().slice(0, 10);
|
||
a.download = `datasets-${projectId}-${formatSuffix}${balanceSuffix}-${dateStr}.${fileExtension}`;
|
||
|
||
document.body.appendChild(a);
|
||
a.click();
|
||
document.body.removeChild(a);
|
||
URL.revokeObjectURL(url);
|
||
};
|
||
|
||
// 导出数据集（保持向后兼容的原有功能）
|
||
const exportDatasets = async exportOptions => {
|
||
try {
|
||
const apiUrl = `/api/projects/${projectId}/datasets/export`;
|
||
const requestBody = {};
|
||
|
||
if (exportOptions.selectedIds && exportOptions.selectedIds.length > 0) {
|
||
requestBody.selectedIds = exportOptions.selectedIds;
|
||
} else if (exportOptions.confirmedOnly) {
|
||
requestBody.status = 'confirmed';
|
||
}
|
||
|
||
if (exportOptions.balanceMode && exportOptions.balanceConfig) {
|
||
requestBody.balanceMode = true;
|
||
requestBody.balanceConfig = exportOptions.balanceConfig;
|
||
}
|
||
|
||
const response = await axios.post(apiUrl, requestBody);
|
||
let dataToExport = response.data;
|
||
|
||
await processAndDownloadData(dataToExport, exportOptions);
|
||
|
||
toast.success(t('datasets.exportSuccess'));
|
||
return true;
|
||
} catch (error) {
|
||
toast.error(error.message);
|
||
return false;
|
||
}
|
||
};
|
||
|
||
// 导出平衡数据集
|
||
const exportBalancedDataset = async exportOptions => {
|
||
const balancedOptions = {
|
||
...exportOptions,
|
||
balanceMode: true,
|
||
balanceConfig: exportOptions.balanceConfig
|
||
};
|
||
return await exportDatasets(balancedOptions);
|
||
};
|
||
|
||
return {
|
||
exportDatasets,
|
||
exportBalancedDataset,
|
||
exportDatasetsStreaming
|
||
};
|
||
};
|
||
|
||
export default useDatasetExport;
|
||
export { useDatasetExport };
|