import IconFont from '@/components/icon-font'; import { BulbOutlined } from '@ant-design/icons'; import { useIntl } from '@umijs/max'; import { Button, Input, Select } from 'antd'; import _ from 'lodash'; import React, { useCallback, useEffect, useRef, useState } from 'react'; import { queryHuggingfaceModels } from '../apis'; import { modelFilesSortOptions, modelSourceMap, ollamaModelOptions } from '../config'; import SearchStyle from '../style/search-result.less'; import SearchInput from './search-input'; import SearchResult from './search-result'; interface SearchInputProps { modelSource: string; onSourceChange?: (source: string) => void; onSelectModel: (model: any) => void; } const sourceList = [ { label: ( ), value: 'huggingface', key: 'huggingface' }, { label: , value: 'ollama_library', key: 'ollama_library' } ]; const SearchModel: React.FC = (props) => { console.log('SearchModel======'); const intl = useIntl(); const { modelSource, onSourceChange, onSelectModel } = props; const [dataSource, setDataSource] = useState<{ repoOptions: any[]; loading: boolean; }>({ repoOptions: [], loading: false }); const [current, setCurrent] = useState(''); const [sortType, setSortType] = useState('downloads'); const cacheRepoOptions = useRef([]); const axiosTokenRef = useRef(null); const customOllamaModelRef = useRef(null); const handleOnSelectModel = useCallback((item: any) => { onSelectModel(item); setCurrent(item.id); }, []); const handleOnSearchRepo = useCallback( async (text: string) => { axiosTokenRef.current?.abort?.(); axiosTokenRef.current = new AbortController(); if (dataSource.loading) return; try { setDataSource((pre) => { pre.loading = true; return { ...pre }; }); cacheRepoOptions.current = []; const params = { search: { query: text, tags: ['gguf'] } }; const models = await queryHuggingfaceModels(params, { signal: axiosTokenRef.current.signal }); const list = _.map(models || [], (item: any) => { return { ...item, value: item.name, label: item.name }; }); const sortedList = _.sortBy( list, (item: any) => item[sortType] ).reverse(); cacheRepoOptions.current = sortedList; setDataSource({ repoOptions: sortedList, loading: false }); handleOnSelectModel(sortedList[0]); } catch (error) { setDataSource({ repoOptions: [], loading: false }); handleOnSelectModel({}); cacheRepoOptions.current = []; } }, [dataSource] ); const handlerSearchModels = useCallback( async (e: any) => { const text = e.target.value; handleOnSearchRepo(text); }, [handleOnSearchRepo] ); const handleOnOpen = () => { if ( !dataSource.repoOptions.length && !cacheRepoOptions.current.length && modelSource === modelSourceMap.huggingface_value ) { handleOnSearchRepo(''); } if (modelSourceMap.ollama_library_value === modelSource) { setDataSource({ repoOptions: ollamaModelOptions, loading: false }); cacheRepoOptions.current = ollamaModelOptions; handleOnSelectModel(ollamaModelOptions[0]); } }; const handleFilterModels = (e: any) => { const text = e.target.value; const list = _.filter(cacheRepoOptions.current, (item: any) => { return item.name.includes(text); }); setDataSource({ repoOptions: list, loading: false }); }; const debounceFilter = _.debounce((e: any) => { handleFilterModels(e); }, 300); const handleSourceChange = (source: string) => { axiosTokenRef.current?.abort?.(); onSourceChange?.(source); setDataSource({ repoOptions: [], loading: false }); cacheRepoOptions.current = []; }; const handleInputChange = (e: any) => { const value = e.target.value; customOllamaModelRef.current = value; }; const handleConfirm = () => { const model = { label: customOllamaModelRef.current, value: customOllamaModelRef.current, name: customOllamaModelRef.current, id: '' }; onSelectModel(model); setCurrent(''); }; const handleSortChange = (value: string) => { const sortedList = _.sortBy( dataSource.repoOptions, (item: any) => item[value] ).reverse(); setSortType(value); setDataSource({ repoOptions: sortedList, loading: false }); }; const renderHFSearch = () => { return ( <> {dataSource.repoOptions.length} results { return ( {intl.formatMessage({ id: 'model.deploy.sort' })}: {label} ); }} options={modelFilesSortOptions} size="middle" style={{ width: '150px' }} > > ); }; const renderOllamaCustom = () => { return ( <> {intl.formatMessage({ id: 'common.button.confirm' })} > ); }; useEffect(() => { handleOnOpen(); console.log('SearchModel useEffect', modelSource); return () => { axiosTokenRef.current?.abort?.(); }; }, [modelSource]); return ( {modelSource === modelSourceMap.huggingface_value ? ( renderHFSearch() ) : ( {intl.formatMessage( { id: 'model.form.ollamatips' }, { name: intl.formatMessage({ id: 'model.form.ollama.model' }) } )} )} { } ); }; export default React.memo(SearchModel);