import useOverlayScroller from '@/hooks/use-overlay-scroller'; import useRequestToken from '@/hooks/use-request-token'; import { ClearOutlined, PlusOutlined, SendOutlined } from '@ant-design/icons'; import { useIntl, useSearchParams } from '@umijs/max'; import { Button, Input, Spin, Tag } from 'antd'; import classNames from 'classnames'; import _ from 'lodash'; import 'overlayscrollbars/overlayscrollbars.css'; import { forwardRef, memo, useCallback, useEffect, useImperativeHandle, useRef, useState } from 'react'; import { rerankerQuery } from '../apis'; import { MessageItem, ParamsSchema } from '../config/types'; import '../style/ground-left.less'; import '../style/rerank.less'; import '../style/system-message-wrap.less'; import DynamicParams from './dynamic-params'; import InputList from './input-list'; import ViewRerankCode from './view-rerank-code'; interface MessageProps { modelList: Global.BaseOption[]; loaded?: boolean; ref?: any; } const paramsConfig: ParamsSchema[] = [ { type: 'InputNumber', name: 'top_n', label: { text: 'Top N', isLocalized: false }, attrs: { min: 1 }, rules: [ { required: true, message: 'Top N is required' } ] } ]; const initialValues = { top_n: 3 }; const GroundReranker: React.FC = forwardRef((props, ref) => { const { modelList } = props; const acceptType = '.txt, .doc, .docx, .xls, .xlsx, .csv, .md, .pdf, .eml, .msg, .ppt, .pptx, .xml, .epub, .html'; const messageId = useRef(0); const [messageList, setMessageList] = useState([]); const intl = useIntl(); const requestSource = useRequestToken(); const [searchParams] = useSearchParams(); const selectModel = searchParams.get('model') || ''; const [parameters, setParams] = useState({}); const [show, setShow] = useState(false); const [loading, setLoading] = useState(false); const [tokenResult, setTokenResult] = useState(null); const [collapse, setCollapse] = useState(false); const contentRef = useRef(''); const scroller = useRef(null); const inputListRef = useRef(null); const paramsRef = useRef(null); const messageListLengthCache = useRef(0); const requestToken = useRef(null); const [fileList, setFileList] = useState< { text: string; name: string; uid: number | string; score?: number; showExtra?: boolean; percent?: number; rank?: number; }[] >([]); const [textList, setTextList] = useState< { text: string; uid: number | string; name: string; score?: number; showExtra?: boolean; percent?: number; rank?: number; }[] >([ { text: '', uid: -1, name: '' }, { text: '', uid: -2, name: '' } ]); const [sortIndexMap, setSortIndexMap] = useState([]); const { initialize, updateScrollerPosition: updateDocumentScrollerPosition } = useOverlayScroller(); const { initialize: innitializeParams, updateScrollerPosition } = useOverlayScroller(); useImperativeHandle(ref, () => { return { viewCode() { setShow(true); }, setCollapse() { setCollapse(!collapse); } }; }); // [0.1, 1.0] const normalizValue = (data: { min: number; max: number; value: number }) => { const range = [0.5, 1.0]; const [a, b] = range; const { min, max, value } = data; if (isNaN(value) || isNaN(min) || isNaN(max) || min > max) { return 0; } if (min === max) { return 100; } const res = a + ((value - min) * (b - a)) / (max - min); return res * 100; }; const renderPercent = useCallback((data: any) => { if (!data.showExtra || !data.percent) { return null; } const percent = data.percent; return (
{intl.formatMessage({ id: 'playground.rerank.rank' })}: {data.rank} {intl.formatMessage({ id: 'playground.rerank.score' })}:{' '} {_.round(data.score, 2)}
); }, []); const setMessageId = () => { messageId.current = messageId.current + 1; return messageId.current; }; const handleStopConversation = () => { requestToken.current?.cancel?.(); setLoading(false); }; const submitMessage = async (current?: { content: string }) => { if (!parameters.model) return; try { setLoading(true); setMessageId(); setTokenResult(null); setSortIndexMap([]); requestToken.current?.cancel?.(); requestToken.current = requestSource(); contentRef.current = current?.content || ''; const documentList: any[] = [...textList, ...fileList]; const result: any = await rerankerQuery( { model: parameters.model, top_n: parameters.top_n, query: contentRef.current, documents: [ ...textList.map((item) => item.text), ...fileList.map((item) => item.text) ] }, { token: requestToken.current.token } ); setMessageId(); setTokenResult(result.usage); const sortList = _.sortBy( result.results || [], (item: any) => item.relevance_score ); const maxValue = sortList[sortList.length - 1].relevance_score; const minValue = sortList[0].relevance_score; // reset state let newTextList = textList.map((item) => { item.percent = undefined; item.score = undefined; item.rank = undefined; return item; }); let sortMap: number[] = []; result.results?.forEach((item: any, sIndex: number) => { sortMap.push(item.index); newTextList[item.index] = { ...newTextList[item.index], uid: setMessageId(), rank: sIndex + 1, score: item.relevance_score, showExtra: true, percent: normalizValue({ min: minValue, max: maxValue, value: item.relevance_score }) }; }); newTextList = _.sortBy(newTextList, 'rank'); setSortIndexMap(sortMap); setTextList(newTextList); setMessageList([ { title: 'Results', role: '', content: result.results?.map((item: any) => { const percent: number = normalizValue({ min: minValue, max: maxValue, value: item.relevance_score }); return { uid: setMessageId(), text: `${item.document?.text?.slice(0, 500) || ''}`, docIndex: item.index, title: documentList[item.index]?.name || '', score: item.relevance_score, extra: renderPercent(percent), normalizValue: percent }; }), uid: setMessageId() } ]); } catch (error: any) { setTokenResult({ error: true, errorMessage: error.response?.data?.error?.message }); } finally { setLoading(false); } }; const handleClear = () => { if (!messageList.length) { return; } setMessageId(); setMessageList([]); setTokenResult(null); }; const handleSendMessage = (message: Omit) => { submitMessage(message); }; const handleSearch = (val: string) => { console.log('val:', val); submitMessage({ content: val }); }; const handleCloseViewCode = () => { setShow(false); }; const handleUpdateFileList = ( files: { text: string; name: string; uid: number | string }[] ) => { console.log('files:', files); setFileList((preList) => { return [...preList, ...files]; }); }; const handleDeleteFile = (uid: number | string) => { setFileList((preList) => { return preList.filter((item) => item.uid !== uid); }); }; const handleAddText = () => { inputListRef.current?.handleAdd(); }; const handleTextListChange = useCallback( (list: { text: string; uid: number | string; name: string }[]) => { const newList = list?.map((item: any) => { item.percent = undefined; item.score = undefined; item.rank = undefined; return item; }); setTextList(newList); }, [] ); const handleOnSort = useCallback( (list: { text: string; uid: number | string; name: string }[]) => { const newList = list?.map((item) => { return { ...item, uid: setMessageId() }; }); setTextList(newList); }, [] ); const handleClearDocuments = () => { setTextList([ { text: '', uid: -1, name: '' }, { text: '', uid: -2, name: '' } ]); setFileList([]); }; useEffect(() => { setMessageId(); setMessageList([]); setTokenResult(null); }, [parameters.model]); useEffect(() => { if (scroller.current) { initialize(scroller.current); } }, [scroller.current, initialize]); useEffect(() => { if (paramsRef.current) { innitializeParams(paramsRef.current); } }, [paramsRef.current, innitializeParams]); useEffect(() => { updateScrollerPosition(); }, [messageList]); useEffect(() => { if (textList.length + fileList.length > messageListLengthCache.current) { updateDocumentScrollerPosition(); } messageListLengthCache.current = textList.length + fileList.length; }, [textList.length, fileList.length]); return (

{intl.formatMessage({ id: 'playground.rerank.query' })}

} placeholder={intl.formatMessage({ id: 'playground.rerank.query.holder' })} >

{intl.formatMessage({ id: 'playground.embedding.documents' })}

{' '} {tokenResult?.total_tokens && ( {intl.formatMessage({ id: 'playground.tokenusage' })}:{' '} {tokenResult?.total_tokens} )}
<>
{loading && (
)}
item.text) .filter((text) => text) }} parameters={{ ...parameters, query: contentRef.current }} onCancel={handleCloseViewCode} title={intl.formatMessage({ id: 'playground.viewcode' })} >
); }); export default memo(GroundReranker);