Spaces:
Running
Running
File size: 2,264 Bytes
e538a38 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 |
import { type ComboboxItem, Select } from "@mantine/core";
import { prebuiltAppConfig } from "@mlc-ai/web-llm";
import { useCallback, useEffect, useState } from "react";
import { isF16Supported } from "../../modules/webGpu";
export default function WebLlmModelSelect({
value,
onChange,
}: {
value: string;
onChange: (value: string) => void;
}) {
const [webGpuModels] = useState<ComboboxItem[]>(() => {
const models = prebuiltAppConfig.model_list
.filter((model) => {
const isSmall = isSmallModel(model);
const suffix = getModelSuffix(isF16Supported, isSmall);
return model.model_id.endsWith(suffix);
})
.sort((a, b) => (a.vram_required_MB ?? 0) - (b.vram_required_MB ?? 0))
.map((model) => {
const modelSizeInMegabytes =
Math.round(model.vram_required_MB ?? 0) || "N/A";
const isSmall = isSmallModel(model);
const suffix = getModelSuffix(isF16Supported, isSmall);
const modelName = model.model_id.replace(suffix, "");
return {
label: `${modelSizeInMegabytes} MB • ${modelName}`,
value: model.model_id,
};
});
return models;
});
useEffect(() => {
const isCurrentModelValid = webGpuModels.some(
(model) => model.value === value,
);
if (!isCurrentModelValid && webGpuModels.length > 0) {
onChange(webGpuModels[0].value);
}
}, [onChange, webGpuModels, value]);
const handleChange = useCallback(
(value: string | null) => {
if (value) onChange(value);
},
[onChange],
);
return (
<Select
value={value}
onChange={handleChange}
label="AI Model"
description="Select the model to use for AI responses."
data={webGpuModels}
allowDeselect={false}
searchable
/>
);
}
type ModelConfig = (typeof prebuiltAppConfig.model_list)[number];
const smallModels = ["SmolLM2-135M", "SmolLM2-360M"] as const;
function isSmallModel(model: ModelConfig) {
return smallModels.some((smallModel) =>
model.model_id.startsWith(smallModel),
);
}
function getModelSuffix(isF16: boolean, isSmall: boolean) {
if (isSmall) return isF16 ? "-q0f16-MLC" : "-q0f32-MLC";
return isF16 ? "-q4f16_1-MLC" : "-q4f32_1-MLC";
}
|