156 lines
4.3 KiB
TypeScript
156 lines
4.3 KiB
TypeScript
import React, { useState, useCallback } from 'react';
|
||
import {
|
||
Box,
|
||
Flex,
|
||
Button,
|
||
Modal,
|
||
ModalOverlay,
|
||
ModalContent,
|
||
ModalHeader,
|
||
ModalCloseButton,
|
||
ModalBody
|
||
} from '@chakra-ui/react';
|
||
import { useToast } from '@/hooks/useToast';
|
||
import { useSelectFile } from '@/hooks/useSelectFile';
|
||
import { customAlphabet } from 'nanoid';
|
||
import { encode } from 'gpt-token-utils';
|
||
import { useConfirm } from '@/hooks/useConfirm';
|
||
import { readTxtContent, readPdfContent, readDocContent } from '@/utils/tools';
|
||
import { useMutation } from '@tanstack/react-query';
|
||
import { postModelDataFileText } from '@/api/model';
|
||
|
||
const nanoid = customAlphabet('abcdefghijklmnopqrstuvwxyz1234567890', 12);
|
||
|
||
const fileExtension = '.txt,.doc,.docx,.pdf,.md';
|
||
|
||
const SelectFileModal = ({
|
||
onClose,
|
||
onSuccess,
|
||
modelId
|
||
}: {
|
||
onClose: () => void;
|
||
onSuccess: () => void;
|
||
modelId: string;
|
||
}) => {
|
||
const [selecting, setSelecting] = useState(false);
|
||
const { toast } = useToast();
|
||
const { File, onOpen } = useSelectFile({ fileType: fileExtension, multiple: true });
|
||
const [fileText, setFileText] = useState('');
|
||
const { openConfirm, ConfirmChild } = useConfirm({
|
||
content: '确认导入该文件，需要一定时间进行拆解，该任务无法终止！如果余额不足，任务讲被终止。'
|
||
});
|
||
|
||
const onSelectFile = useCallback(
|
||
async (e: File[]) => {
|
||
setSelecting(true);
|
||
try {
|
||
const fileTexts = (
|
||
await Promise.all(
|
||
e.map((file) => {
|
||
// @ts-ignore
|
||
const extension = file?.name?.split('.').pop().toLowerCase();
|
||
switch (extension) {
|
||
case 'txt':
|
||
case 'md':
|
||
return readTxtContent(file);
|
||
case 'pdf':
|
||
return readPdfContent(file);
|
||
case 'doc':
|
||
case 'docx':
|
||
return readDocContent(file);
|
||
default:
|
||
return '';
|
||
}
|
||
})
|
||
)
|
||
)
|
||
.join('\n')
|
||
.replace(/\n+/g, '\n');
|
||
setFileText(fileTexts);
|
||
console.log(encode(fileTexts));
|
||
} catch (error: any) {
|
||
console.log(error);
|
||
toast({
|
||
title: typeof error === 'string' ? error : '解析文件失败',
|
||
status: 'error'
|
||
});
|
||
}
|
||
setSelecting(false);
|
||
},
|
||
[setSelecting, toast]
|
||
);
|
||
|
||
const { mutate, isLoading } = useMutation({
|
||
mutationFn: async () => {
|
||
if (!fileText) return;
|
||
await postModelDataFileText(modelId, fileText);
|
||
toast({
|
||
title: '导入数据成功,需要一段拆解和训练',
|
||
status: 'success'
|
||
});
|
||
onClose();
|
||
onSuccess();
|
||
},
|
||
onError() {
|
||
toast({
|
||
title: '导入文件失败',
|
||
status: 'error'
|
||
});
|
||
}
|
||
});
|
||
|
||
return (
|
||
<Modal isOpen={true} onClose={onClose}>
|
||
<ModalOverlay />
|
||
<ModalContent maxW={'min(900px, 90vw)'} position={'relative'}>
|
||
<ModalHeader>文件导入</ModalHeader>
|
||
<ModalCloseButton />
|
||
|
||
<ModalBody>
|
||
<Flex
|
||
flexDirection={'column'}
|
||
p={2}
|
||
h={'100%'}
|
||
alignItems={'center'}
|
||
justifyContent={'center'}
|
||
fontSize={'sm'}
|
||
>
|
||
<Button isLoading={selecting} onClick={onOpen}>
|
||
选择文件
|
||
</Button>
|
||
<Box mt={2}>支持 {fileExtension} 文件. 会先对文本进行拆分，需要时间较长。</Box>
|
||
<Box mt={2}>
|
||
一共 {fileText.length} 个字，{encode(fileText).length} 个tokens
|
||
</Box>
|
||
<Box
|
||
h={'300px'}
|
||
w={'100%'}
|
||
overflow={'auto'}
|
||
p={2}
|
||
backgroundColor={'blackAlpha.50'}
|
||
whiteSpace={'pre'}
|
||
fontSize={'xs'}
|
||
>
|
||
{fileText}
|
||
</Box>
|
||
</Flex>
|
||
</ModalBody>
|
||
|
||
<Flex px={6} pt={2} pb={4}>
|
||
<Box flex={1}></Box>
|
||
<Button variant={'outline'} mr={3} onClick={onClose}>
|
||
取消
|
||
</Button>
|
||
<Button isLoading={isLoading} isDisabled={fileText === ''} onClick={openConfirm(mutate)}>
|
||
确认导入
|
||
</Button>
|
||
</Flex>
|
||
</ModalContent>
|
||
<ConfirmChild />
|
||
<File onSelect={onSelectFile} />
|
||
</Modal>
|
||
);
|
||
};
|
||
|
||
export default SelectFileModal;
|