606 lines
16 KiB
TypeScript
606 lines
16 KiB
TypeScript
import { getRequestId, setRquestId } from '@/atoms/models';
|
|
import BaseSelect from '@/components/seal-form/base/select';
|
|
import { createAxiosToken } from '@/hooks/use-chunk-request';
|
|
import ColumnWrapper from '@/pages/_components/column-wrapper';
|
|
import { useIntl } from '@umijs/max';
|
|
import { Pagination } from 'antd';
|
|
import _ from 'lodash';
|
|
import React, { useEffect, useRef, useState } from 'react';
|
|
import styled from 'styled-components';
|
|
import {
|
|
evaluationsModelSpec,
|
|
queryHuggingfaceModels,
|
|
queryModelScopeModels
|
|
} from '../apis';
|
|
import { ModelScopeSortType, ModelSortType, modelSourceMap } from '../config';
|
|
import { MessageStatus, WarningStausOptions } from '../hooks';
|
|
import useCheckBackend from '../hooks/use-check-backend';
|
|
import useRecognizeAudio from '../hooks/use-recognize-audio';
|
|
import SearchStyle from '../style/search-result.less';
|
|
import SearchInput from './search-input';
|
|
import SearchResult from './search-result';
|
|
|
|
const filterOptions = [
|
|
{ label: 'FP8', value: 'fp8' },
|
|
{ label: 'AWQ', value: 'awq' },
|
|
{ label: 'GPTQ', value: 'gptq' }
|
|
];
|
|
const PaginationMain = styled(Pagination)`
|
|
.ant-pagination-slash {
|
|
margin-inline: 5px;
|
|
}
|
|
`;
|
|
|
|
interface SearchInputProps {
|
|
hasLinuxWorker?: boolean;
|
|
modelSource: string;
|
|
isDownload?: boolean;
|
|
gpuOptions?: any[];
|
|
clusterId?: number;
|
|
setLoadingModel?: (flag: boolean) => void;
|
|
onSourceChange?: (source: string) => void;
|
|
onSelectModel: (model: any, manul?: boolean) => void;
|
|
onSelectModelAfterEvaluate?: (model: any, manual?: boolean) => void;
|
|
displayEvaluateStatus?: (
|
|
data: MessageStatus,
|
|
options?: WarningStausOptions
|
|
) => void;
|
|
}
|
|
|
|
const SearchModel: React.FC<SearchInputProps> = (props) => {
|
|
const intl = useIntl();
|
|
const {
|
|
modelSource,
|
|
isDownload,
|
|
gpuOptions,
|
|
clusterId,
|
|
setLoadingModel,
|
|
onSelectModel,
|
|
onSelectModelAfterEvaluate,
|
|
displayEvaluateStatus
|
|
} = props;
|
|
|
|
const { recognizeAudioModel } = useRecognizeAudio();
|
|
const { checkCurrentbackend } = useCheckBackend();
|
|
const [dataSource, setDataSource] = useState<{
|
|
dataList: any[];
|
|
loading: boolean;
|
|
networkError: boolean;
|
|
sortType: string;
|
|
filters: Record<string, any>;
|
|
}>({
|
|
dataList: [],
|
|
loading: false,
|
|
networkError: false,
|
|
sortType: ModelSortType.trendingScore,
|
|
filters: {
|
|
tag: null
|
|
}
|
|
});
|
|
const SUPPORTEDSOURCE = [
|
|
modelSourceMap.huggingface_value,
|
|
modelSourceMap.modelscope_value
|
|
];
|
|
|
|
const [isEvaluating, setIsEvaluating] = useState<boolean>(false);
|
|
const [current, setCurrent] = useState<string>('');
|
|
const currentRef = useRef<string>('');
|
|
const cacheRepoOptions = useRef<any[]>([]);
|
|
const axiosTokenRef = useRef<any>(null);
|
|
const checkTokenRef = useRef<any>(null);
|
|
const searchInputRef = useRef<any>('');
|
|
const timer = useRef<any>(null);
|
|
const requestIdRef = useRef<number>(0);
|
|
const searchRepoRequestIdRef = useRef<number>(0);
|
|
const [paginationInfo, setPaginationInfo] = useState({
|
|
page: 1,
|
|
perPage: 10,
|
|
total: 0
|
|
});
|
|
const modelFilesSortOptions = [
|
|
{
|
|
label: intl.formatMessage({ id: 'models.sort.trending' }),
|
|
value: ModelSortType.trendingScore
|
|
},
|
|
{
|
|
label: intl.formatMessage({ id: 'models.sort.likes' }),
|
|
value: ModelSortType.likes
|
|
},
|
|
{
|
|
label: intl.formatMessage({ id: 'models.sort.downloads' }),
|
|
value: ModelSortType.downloads
|
|
},
|
|
{
|
|
label: intl.formatMessage({ id: 'models.sort.updated' }),
|
|
value: ModelSortType.lastModified
|
|
}
|
|
];
|
|
|
|
const updateSearchRepoRequestId = () => {
|
|
searchRepoRequestIdRef.current += 1;
|
|
return searchRepoRequestIdRef.current;
|
|
};
|
|
|
|
const updateRequestId = () => {
|
|
requestIdRef.current += 1;
|
|
return requestIdRef.current;
|
|
};
|
|
|
|
const checkIsGGUF = (item: any) => {
|
|
const isGGUF = _.some(item.tags, (tag: string) => {
|
|
return tag.toLowerCase() === 'gguf';
|
|
});
|
|
const isGGUFFromMs = _.some(item.libraries, (tag: string) => {
|
|
return tag.toLowerCase() === 'gguf';
|
|
});
|
|
return isGGUF || isGGUFFromMs;
|
|
};
|
|
|
|
const handleOnSelectModel = (model: any, manual?: boolean) => {
|
|
const item = model || {};
|
|
if (!item.evaluated || item.isGGUF) {
|
|
onSelectModel(item, manual);
|
|
} else {
|
|
onSelectModelAfterEvaluate?.(item, manual);
|
|
}
|
|
setCurrent(item.id);
|
|
currentRef.current = item.id;
|
|
};
|
|
|
|
// huggeface
|
|
const getModelsFromHuggingface = async (query: {
|
|
sort: string;
|
|
filters?: Record<string, any>;
|
|
}) => {
|
|
const currentSearchId = setRquestId();
|
|
const task: any = searchInputRef.current ? '' : 'text-generation';
|
|
const params = {
|
|
search: {
|
|
query: searchInputRef.current || '',
|
|
sort: query.sort,
|
|
tags: query.filters?.tag ? [query.filters.tag] : [],
|
|
task: task
|
|
}
|
|
};
|
|
const data = await queryHuggingfaceModels(params, {
|
|
signal: axiosTokenRef.current.signal
|
|
});
|
|
if (getRequestId() !== currentSearchId) {
|
|
throw 'new request has been sent';
|
|
}
|
|
let list = _.map(data || [], (item: any) => {
|
|
return {
|
|
...item,
|
|
value: item.name,
|
|
label: item.name,
|
|
isGGUF: checkIsGGUF(item),
|
|
source: modelSource
|
|
};
|
|
});
|
|
return list;
|
|
};
|
|
|
|
// modelscope, only modelscope has page and perPage
|
|
const getModelsFromModelscope = async (queryParams: {
|
|
sortType: string;
|
|
page: number;
|
|
perPage?: number;
|
|
filters?: Record<string, any>;
|
|
}) => {
|
|
const currentSearchId = setRquestId();
|
|
try {
|
|
const params = {
|
|
Name: `${searchInputRef.current}`,
|
|
tags: queryParams.filters?.tag ? [queryParams.filters.tag] : [],
|
|
SortBy: ModelScopeSortType[queryParams.sortType],
|
|
PageNumber: queryParams.page,
|
|
PageSize: queryParams.perPage,
|
|
tasks: []
|
|
};
|
|
const data = await queryModelScopeModels(params, {
|
|
signal: axiosTokenRef.current.signal
|
|
});
|
|
|
|
if (getRequestId() !== currentSearchId) {
|
|
throw 'new request has been sent';
|
|
}
|
|
|
|
let list = _.map(_.get(data, 'Data.Model.Models') || [], (item: any) => {
|
|
return {
|
|
path: item.Path,
|
|
name: `${item.Path}/${item.Name}`,
|
|
downloads: item.Downloads,
|
|
id: `${item.Path}/${item.Name}`,
|
|
updatedAt: item.LastUpdatedTime * 1000,
|
|
likes: item.Stars,
|
|
value: item.Name,
|
|
label: item.Name,
|
|
revision: item.Revision,
|
|
task: item.Tasks?.map((sItem: any) => sItem.Name).join(','),
|
|
tags: item.Tags,
|
|
libraries: item.Libraries,
|
|
avatar: item.Avatar,
|
|
isGGUF: checkIsGGUF({
|
|
tags: item.Tags,
|
|
libraries: item.Libraries
|
|
}),
|
|
source: modelSource
|
|
};
|
|
});
|
|
|
|
setPaginationInfo((prev) => {
|
|
return {
|
|
...prev,
|
|
page: queryParams.page,
|
|
total: _.get(data, 'Data.Model.TotalCount', 0)
|
|
};
|
|
});
|
|
return list;
|
|
} catch (error) {
|
|
setPaginationInfo((prev) => {
|
|
return {
|
|
...prev,
|
|
page: queryParams.page
|
|
};
|
|
});
|
|
throw error;
|
|
}
|
|
};
|
|
|
|
const getEvaluateResults = async (repoList: any[]) => {
|
|
checkTokenRef.current?.cancel?.();
|
|
checkTokenRef.current = createAxiosToken();
|
|
const evaluations = await evaluationsModelSpec(
|
|
{
|
|
cluster_id: clusterId!,
|
|
model_specs: repoList
|
|
},
|
|
{
|
|
token: checkTokenRef.current?.token
|
|
}
|
|
);
|
|
return evaluations.results;
|
|
};
|
|
|
|
const handleEvaluate = async (list: any[]) => {
|
|
if (isDownload) {
|
|
return;
|
|
}
|
|
const currentRequestId = updateRequestId();
|
|
const currentSearchId = getRequestId();
|
|
try {
|
|
const repoList = list.map((item) => {
|
|
const res = recognizeAudioModel(item, modelSource);
|
|
|
|
let backendObj = {};
|
|
const backend = checkCurrentbackend({
|
|
isGGUF: item.isGGUF,
|
|
isAudio: res.isAudio,
|
|
gpuOptions: gpuOptions || []
|
|
});
|
|
|
|
if (backend) {
|
|
backendObj = {
|
|
backend: backend
|
|
};
|
|
}
|
|
|
|
return {
|
|
...backendObj,
|
|
cluster_id: clusterId,
|
|
source: modelSource,
|
|
...(modelSource === modelSourceMap.huggingface_value
|
|
? {
|
|
huggingface_repo_id: item.name
|
|
}
|
|
: {
|
|
model_scope_model_id: item.name
|
|
})
|
|
};
|
|
});
|
|
|
|
setIsEvaluating(true);
|
|
const evaluations = await getEvaluateResults(repoList);
|
|
// bind the requestId to the current request and searchId
|
|
if (
|
|
requestIdRef.current !== currentRequestId &&
|
|
currentSearchId !== getRequestId()
|
|
) {
|
|
return;
|
|
}
|
|
const resultList = list.map((item, index) => {
|
|
return {
|
|
...item,
|
|
evaluated: true,
|
|
evaluateResult: evaluations[index] || null
|
|
};
|
|
});
|
|
setIsEvaluating(false);
|
|
|
|
setDataSource((pre) => {
|
|
return {
|
|
...pre,
|
|
loading: false,
|
|
dataList: resultList
|
|
};
|
|
});
|
|
// current selected item
|
|
const currentItem = resultList.find(
|
|
(item) => item.id === currentRef.current
|
|
);
|
|
|
|
// if it is gguf, would trigger a evaluation after select a model file
|
|
if (currentItem && !currentItem.isGGUF) {
|
|
onSelectModelAfterEvaluate?.(currentItem);
|
|
}
|
|
} catch (error) {
|
|
// cancel the corrponding request
|
|
if (requestIdRef.current === currentRequestId) {
|
|
setIsEvaluating(false);
|
|
}
|
|
}
|
|
};
|
|
|
|
const getCurrentPage = (page: number) => {
|
|
const start = (page - 1) * paginationInfo.perPage;
|
|
const end = start + paginationInfo.perPage;
|
|
return cacheRepoOptions.current.slice(start, end);
|
|
};
|
|
|
|
const handleOnSearchRepo = async (params: {
|
|
sortType: string;
|
|
page: number;
|
|
perPage: number;
|
|
filters: Record<string, any>;
|
|
}) => {
|
|
if (!SUPPORTEDSOURCE.includes(modelSource)) {
|
|
return;
|
|
}
|
|
const currentSearchId = updateSearchRepoRequestId();
|
|
axiosTokenRef.current?.abort?.('cancel previous request');
|
|
axiosTokenRef.current = new AbortController();
|
|
checkTokenRef.current?.cancel?.();
|
|
if (timer.current) {
|
|
clearTimeout(timer.current);
|
|
}
|
|
try {
|
|
setDataSource((pre) => {
|
|
pre.loading = true;
|
|
return { ...pre };
|
|
});
|
|
setLoadingModel?.(true);
|
|
cacheRepoOptions.current = [];
|
|
let list: any[] = [];
|
|
if (modelSource === modelSourceMap.huggingface_value) {
|
|
const resultList = await getModelsFromHuggingface({
|
|
sort: params.sortType,
|
|
filters: params.filters
|
|
});
|
|
|
|
cacheRepoOptions.current = resultList;
|
|
|
|
// hf has no page and perPage, so we need to slice the resultList
|
|
list = getCurrentPage(params.page);
|
|
setPaginationInfo((prev) => {
|
|
return {
|
|
...prev,
|
|
page: params.page,
|
|
total: resultList.length
|
|
};
|
|
});
|
|
} else if (modelSource === modelSourceMap.modelscope_value) {
|
|
list = await getModelsFromModelscope(params);
|
|
console.log('list:', list);
|
|
cacheRepoOptions.current = list;
|
|
}
|
|
|
|
setDataSource({
|
|
dataList: list,
|
|
loading: false,
|
|
networkError: false,
|
|
sortType: params.sortType,
|
|
filters: params.filters
|
|
});
|
|
|
|
handleOnSelectModel(list[0]);
|
|
setLoadingModel?.(false);
|
|
|
|
handleEvaluate(list);
|
|
} catch (error: any) {
|
|
console.log('error:', error);
|
|
setDataSource({
|
|
dataList: [],
|
|
loading: currentSearchId !== searchRepoRequestIdRef.current,
|
|
sortType: params.sortType,
|
|
filters: params.filters,
|
|
networkError: error?.message === 'Failed to fetch'
|
|
});
|
|
|
|
setLoadingModel?.(currentSearchId !== searchRepoRequestIdRef.current);
|
|
displayEvaluateStatus?.({
|
|
show: false,
|
|
message: ''
|
|
});
|
|
handleOnSelectModel({});
|
|
cacheRepoOptions.current = [];
|
|
}
|
|
};
|
|
const handleSearchInputChange = (e: any) => {
|
|
searchInputRef.current = e.target.value;
|
|
};
|
|
const handlerSearchModels = _.debounce(
|
|
() =>
|
|
handleOnSearchRepo({
|
|
sortType: dataSource.sortType,
|
|
page: 1,
|
|
perPage: paginationInfo.perPage,
|
|
filters: dataSource.filters
|
|
}),
|
|
100
|
|
);
|
|
|
|
const handleOnOpen = () => {
|
|
if (
|
|
!dataSource.dataList.length &&
|
|
!cacheRepoOptions.current.length &&
|
|
SUPPORTEDSOURCE.includes(modelSource)
|
|
) {
|
|
handleOnSearchRepo({
|
|
sortType: dataSource.sortType,
|
|
page: 1,
|
|
perPage: paginationInfo.perPage,
|
|
filters: dataSource.filters
|
|
});
|
|
}
|
|
};
|
|
|
|
const handleSortChange = (value: string) => {
|
|
handleOnSearchRepo({
|
|
sortType: value,
|
|
page: 1,
|
|
perPage: paginationInfo.perPage,
|
|
filters: dataSource.filters
|
|
});
|
|
};
|
|
|
|
const handleFilterChange = (value: string) => {
|
|
handleOnSearchRepo({
|
|
sortType: dataSource.sortType,
|
|
page: 1,
|
|
perPage: paginationInfo.perPage,
|
|
filters: {
|
|
...dataSource.filters,
|
|
tag: value
|
|
}
|
|
});
|
|
};
|
|
|
|
const handleOnPageChange = (page: number) => {
|
|
if (modelSource === modelSourceMap.huggingface_value) {
|
|
const currentList = getCurrentPage(page);
|
|
setPaginationInfo((prev) => {
|
|
return {
|
|
...prev,
|
|
page: page
|
|
};
|
|
});
|
|
setDataSource((pre) => {
|
|
return {
|
|
...pre,
|
|
dataList: currentList
|
|
};
|
|
});
|
|
handleOnSelectModel(currentList[0]);
|
|
handleEvaluate(currentList);
|
|
} else if (modelSource === modelSourceMap.modelscope_value) {
|
|
setPaginationInfo((prev) => {
|
|
return {
|
|
...prev,
|
|
page: page
|
|
};
|
|
});
|
|
handleOnSearchRepo({
|
|
sortType: dataSource.sortType,
|
|
page: page,
|
|
perPage: paginationInfo.perPage,
|
|
filters: dataSource.filters
|
|
});
|
|
}
|
|
};
|
|
|
|
const handleSelectModelManually = (model: any) => {
|
|
if (model.id === currentRef.current) {
|
|
return;
|
|
}
|
|
setRquestId();
|
|
handleOnSelectModel(model, true);
|
|
};
|
|
|
|
const renderHFSearch = () => {
|
|
return (
|
|
<>
|
|
<SearchInput
|
|
onSearch={handlerSearchModels}
|
|
onChange={handleSearchInputChange}
|
|
modelSource={modelSource}
|
|
></SearchInput>
|
|
<div className={SearchStyle.filter}>
|
|
<span className="flex-center gap-8">
|
|
<BaseSelect
|
|
value={dataSource.sortType}
|
|
onChange={handleSortChange}
|
|
prefix={
|
|
<span>{intl.formatMessage({ id: 'model.deploy.sort' })}:</span>
|
|
}
|
|
options={modelFilesSortOptions}
|
|
size="middle"
|
|
style={{ width: '150px' }}
|
|
></BaseSelect>
|
|
<BaseSelect
|
|
allowClear
|
|
value={dataSource.filters.tag}
|
|
onChange={handleFilterChange}
|
|
options={filterOptions}
|
|
size="middle"
|
|
placeholder={intl.formatMessage({
|
|
id: 'models.form.quantization'
|
|
})}
|
|
style={{ width: 130 }}
|
|
></BaseSelect>
|
|
</span>
|
|
<PaginationMain
|
|
simple={{ readOnly: true }}
|
|
total={paginationInfo.total}
|
|
current={paginationInfo.page}
|
|
pageSize={paginationInfo.perPage}
|
|
onChange={handleOnPageChange}
|
|
showSizeChanger={false}
|
|
hideOnSinglePage={paginationInfo.total <= paginationInfo.perPage}
|
|
></PaginationMain>
|
|
</div>
|
|
</>
|
|
);
|
|
};
|
|
|
|
useEffect(() => {
|
|
handleOnOpen();
|
|
}, [modelSource]);
|
|
|
|
useEffect(() => {
|
|
return () => {
|
|
axiosTokenRef.current?.abort?.();
|
|
checkTokenRef.current?.cancel?.();
|
|
|
|
if (timer.current) {
|
|
clearTimeout(timer.current);
|
|
}
|
|
};
|
|
}, []);
|
|
|
|
return (
|
|
<div style={{ width: '100%' }}>
|
|
<div className={SearchStyle['search-bar']}>{renderHFSearch()}</div>
|
|
<ColumnWrapper
|
|
maxHeight={'calc(100vh - 210px)'}
|
|
styles={{
|
|
container: {
|
|
paddingTop: 0
|
|
}
|
|
}}
|
|
>
|
|
<SearchResult
|
|
loading={dataSource.loading}
|
|
resultList={dataSource.dataList}
|
|
networkError={dataSource.networkError}
|
|
current={current}
|
|
source={modelSource}
|
|
isEvaluating={isEvaluating}
|
|
onSelect={handleSelectModelManually}
|
|
></SearchResult>
|
|
</ColumnWrapper>
|
|
</div>
|
|
);
|
|
};
|
|
|
|
export default SearchModel;
|