jan/web/screens/HubScreen2/components/ModelSearchBar.tsx
2024-08-06 13:17:38 +07:00

114 lines
3.4 KiB
TypeScript

import React, { useCallback, useState } from 'react'
import { Button, Input } from '@janhq/joi'
import { useSetAtom } from 'jotai'
import { SearchIcon } from 'lucide-react'
import { FoldersIcon } from 'lucide-react'
import { useDebouncedCallback } from 'use-debounce'
import { toaster } from '@/containers/Toast'
import { useGetHFRepoData } from '@/hooks/useGetHFRepoData'
// import { setImportModelStageAtom } from '@/hooks/useImportModel'
import { MainViewState, mainViewStateAtom } from '@/helpers/atoms/App.atom'
import {
importHuggingFaceModelStageAtom,
importingHuggingFaceRepoDataAtom,
} from '@/helpers/atoms/HuggingFace.atom'
import { selectedSettingAtom } from '@/helpers/atoms/Setting.atom'
type Props = {
onSearchChanged: (query: string) => void
}
const ModelSearchBar: React.FC<Props> = ({ onSearchChanged }) => {
const [searchText, setSearchText] = useState('')
const { getHfRepoData } = useGetHFRepoData()
const setMainViewState = useSetAtom(mainViewStateAtom)
const setSelectedSetting = useSetAtom(selectedSettingAtom)
// const setImportModelStage = useSetAtom(setImportModelStageAtom)
const setImportingHuggingFaceRepoData = useSetAtom(
importingHuggingFaceRepoDataAtom
)
const setImportHuggingFaceModelStage = useSetAtom(
importHuggingFaceModelStageAtom
)
// const onImportModelClick = useCallback(() => {
// setImportModelStage('SELECTING_MODEL')
// }, [setImportModelStage])
const debounced = useDebouncedCallback(async (searchText: string) => {
if (searchText.indexOf('/') === -1) {
// If we don't find / in the text, perform a local search
onSearchChanged?.(searchText)
return
}
try {
const data = await getHfRepoData(searchText)
setImportingHuggingFaceRepoData(data)
setImportHuggingFaceModelStage('REPO_DETAIL')
} catch (err) {
let errMessage = 'Unexpected Error'
if (err instanceof Error) {
errMessage = err.message
}
toaster({
title: 'Failed to get Hugging Face models',
description: errMessage,
type: 'error',
})
console.error(err)
}
}, 300)
const onQueryChanged = useCallback(
(e: React.ChangeEvent<HTMLInputElement>) => {
e.preventDefault()
e.stopPropagation()
const text = e.target.value
setSearchText(text)
debounced(text)
},
[debounced]
)
return (
<div className="mx-4 mt-4 flex h-[128px] items-center justify-center gap-3 rounded-[10px] bg-blue-400">
<Input
className="w-full bg-[hsla(var(--app-bg))] md:w-[320px]"
prefixIcon={<SearchIcon size={16} />}
placeholder="Search or paste Hugging Face URL"
value={searchText}
onChange={onQueryChanged}
/>
<Button
className="flex gap-2 bg-[hsla(var(--app-bg))] text-[hsla(var(--text-primary))]"
theme="ghost"
onClick={() => {
setMainViewState(MainViewState.Settings)
setSelectedSetting('My Models')
}}
>
<FoldersIcon size={16} />
<span>My models</span>
</Button>
{/* Temporary hidden button import model */}
{/* <Button
className="flex gap-2 bg-[hsla(var(--app-bg))] text-[hsla(var(--text-primary))]"
theme="ghost"
onClick={onImportModelClick}
>
<ImportIcon size={16} />
<span>Import model</span>
</Button> */}
</div>
)
}
export default ModelSearchBar