248 lines
7.4 KiB
TypeScript
248 lines
7.4 KiB
TypeScript
import { useEffect, useRef, useState } from "react";
|
|
import SuggestionBox from "./suggestionBox";
|
|
import { queryAtom } from "lib/state/query";
|
|
import { suggestionItem, suggestionsResponse } from "global";
|
|
import getSearchEngineName from "lib/onesearch/getSearchEngineName";
|
|
import PlainSearch from "./plainSearch";
|
|
import { suggestionAtom } from "lib/state/suggestion";
|
|
import validLink from "lib/url/validLink";
|
|
import LinkSuggestion from "./link";
|
|
import { selectedSuggestionAtom } from "lib/state/suggestionSelection";
|
|
import { settingsAtom } from "lib/state/settings";
|
|
import PlainText from "./plainText";
|
|
import { sendError } from "lib/telemetering/sendError";
|
|
import { handleNLUResult } from "./handleNLUResult";
|
|
import * as ort from 'onnxruntime-web';
|
|
import { useAtom, useAtomValue } from "jotai";
|
|
import i18next from "i18next";
|
|
import { useTranslation } from "react-i18next";
|
|
import { keywordSuggestion } from "lib/onesearch/keywordSuggestion";
|
|
import tokenize from "lib/nlp/tokenizer";
|
|
import { getEmbedding, getEmbeddingLayer } from "lib/nlp/getEmbedding";
|
|
|
|
interface EmbeddingLayer {
|
|
[key: number]: Float32Array<ArrayBufferLike>;
|
|
}
|
|
|
|
export default function OneSearch() {
|
|
const [suggestion, setFinalSuggetsion] = useAtom(suggestionAtom);
|
|
const [embeddingLayer, setEmbeddingLayer] = useState<EmbeddingLayer | null>(null);
|
|
const [NLUsession, setNLUsession] = useState<ort.InferenceSession | null>(null);
|
|
const lastRequestTimeRef = useRef(0);
|
|
const selected = useAtomValue(selectedSuggestionAtom);
|
|
const settings = useAtomValue(settingsAtom);
|
|
const devMode = false;
|
|
const query = useAtomValue(queryAtom);
|
|
const engineName = getSearchEngineName();
|
|
const engine = settings.currentSearchEngine;
|
|
const { t } = useTranslation();
|
|
const lang = i18next.language;
|
|
|
|
useEffect(() => {
|
|
const time = new Date().getTime().toString();
|
|
if (query.trim() === "" || query.length > 120) {
|
|
cleanSuggestion("QUERY", "NAVIGATION");
|
|
return;
|
|
}
|
|
fetch(`/api/v1/suggestion?q=${query}&l=${lang}&t=${time}&engine=${engine}`)
|
|
.then((res) => res.json())
|
|
.then((data: suggestionsResponse) => {
|
|
try {
|
|
const suggestionToUpdate: suggestionItem[] = data.suggestions;
|
|
if (data.time > lastRequestTimeRef.current) {
|
|
cleanSuggestion("NAVIGATION", "QUERY");
|
|
lastRequestTimeRef.current = data.time;
|
|
updateSuggestion(suggestionToUpdate);
|
|
}
|
|
} catch (error: Error | unknown) {
|
|
if (error instanceof Error) {
|
|
sendError(error);
|
|
}
|
|
}
|
|
})
|
|
.catch((error) => {
|
|
// Handle fetch error
|
|
sendError(error);
|
|
});
|
|
}, [query]);
|
|
|
|
function updateSuggestion(data: suggestionItem[]) {
|
|
setFinalSuggetsion((cur: suggestionItem[]) => {
|
|
const types: string[] = [];
|
|
for (const sug of data) {
|
|
if (!types.includes(sug.type)) types.push(sug.type);
|
|
}
|
|
for (const type of types) {
|
|
cur = cur.filter((item) => {
|
|
return item.type !== type;
|
|
});
|
|
}
|
|
return cur.concat(data).sort((a, b) => {
|
|
return b.relevance - a.relevance;
|
|
});
|
|
});
|
|
}
|
|
|
|
function cleanSuggestion(...types: string[]) {
|
|
setFinalSuggetsion((suggestion: suggestionItem[]) => {
|
|
return suggestion.filter((item) => {
|
|
return !types.includes(item.type);
|
|
});
|
|
});
|
|
}
|
|
|
|
useEffect(() => {
|
|
if (embeddingLayer !== null) return;
|
|
const embedding_file = "/models/token_embeddings.bin";
|
|
(async function () {
|
|
const result = await fetch(embedding_file);
|
|
const arrBuf = await result.arrayBuffer();
|
|
const embeddingDict = getEmbeddingLayer(arrBuf);
|
|
setEmbeddingLayer(embeddingDict);
|
|
|
|
await loadModel("/models/NLU.onnx");
|
|
// if (!modelLoaded) {
|
|
// console.error("NLU model was not correctly loaded.")
|
|
// }
|
|
})();
|
|
|
|
}, []);
|
|
|
|
async function loadModel(modelPath: string) {
|
|
ort.env.wasm.wasmPaths = "/onnx/"
|
|
const session = await ort.InferenceSession.create(modelPath);
|
|
setNLUsession(session);
|
|
}
|
|
|
|
async function getNLUResult(query: string) {
|
|
const start = new Date().getTime();
|
|
if (embeddingLayer === null || NLUsession === null) return;
|
|
const tokenIds = await tokenize(query, "Qwen/Qwen2.5-3B");
|
|
console.log(new Date().getTime() - start, "ms");
|
|
const embeddings = getEmbedding(tokenIds, embeddingLayer, 64);
|
|
const inputTensor = new ort.Tensor("float32", embeddings, [1, 64, 96]);
|
|
const feeds = { input: inputTensor };
|
|
console.log(new Date().getTime() - start, "ms");
|
|
const results = await NLUsession.run(feeds);
|
|
console.log(new Date().getTime() - start, "ms");
|
|
return results;
|
|
}
|
|
|
|
useEffect(() => {
|
|
cleanSuggestion("default-link", "default", "text", "link");
|
|
if (validLink(query)) {
|
|
updateSuggestion([
|
|
{
|
|
type: "default-link",
|
|
suggestion: query,
|
|
relevance: 3000,
|
|
prompt: <span>Go to: </span>
|
|
},
|
|
{ type: "default", suggestion: query, relevance: 1600 }
|
|
]);
|
|
} else {
|
|
updateSuggestion([
|
|
{
|
|
type: "default",
|
|
suggestion: query,
|
|
relevance: 2000
|
|
}
|
|
]);
|
|
}
|
|
|
|
if (keywordSuggestion(query) !== null) {
|
|
updateSuggestion([keywordSuggestion(query)!]);
|
|
}
|
|
|
|
(async function () {
|
|
const result = await getNLUResult(query);
|
|
console.log(result);
|
|
})();
|
|
}, [query, engineName]);
|
|
|
|
return (
|
|
<SuggestionBox>
|
|
{suggestion.map((s, i) => {
|
|
if (s.suggestion.trim() === "") return;
|
|
if (s.type === "default") {
|
|
return (
|
|
<PlainSearch key={i} query={s.suggestion} selected={i == selected}>
|
|
{s.suggestion}
|
|
<span className="text-zinc-700 dark:text-zinc-400 text-sm">
|
|
{t("search.search-help-text", { engine: engineName })}
|
|
</span>
|
|
{devMode && (
|
|
<span className="absolute text-zinc-700 dark:text-zinc-400 text-sm leading-10 h-10 right-2">
|
|
{s.relevance}
|
|
</span>
|
|
)}
|
|
</PlainSearch>
|
|
);
|
|
} else if (s.type === "QUERY") {
|
|
return (
|
|
<PlainSearch key={i} query={s.suggestion} selected={i == selected}>
|
|
{s.suggestion}
|
|
{devMode && (
|
|
<span className="absolute text-zinc-700 dark:text-zinc-400 text-sm leading-10 h-10 right-2">
|
|
{s.relevance}
|
|
</span>
|
|
)}
|
|
</PlainSearch>
|
|
);
|
|
} else if (
|
|
s.type === "NAVIGATION" ||
|
|
s.type === "default-link" ||
|
|
s.type === "link"
|
|
) {
|
|
return (
|
|
<LinkSuggestion key={i} query={s.suggestion} selected={i == selected}>
|
|
{s.prompt && (
|
|
<span className="text-zinc-700 dark:text-zinc-400">{s.prompt}</span>
|
|
)}
|
|
{s.suggestion}
|
|
{devMode && (
|
|
<span className="absolute text-zinc-700 dark:text-zinc-400 text-sm leading-10 h-10 right-2">
|
|
{s.relevance}
|
|
</span>
|
|
)}
|
|
</LinkSuggestion>
|
|
);
|
|
} else if (s.type === "text") {
|
|
return (
|
|
<PlainText key={i} selected={i == selected}>
|
|
{s.prompt && (
|
|
<span className="text-zinc-700 dark:text-zinc-400">{s.prompt}</span>
|
|
)}
|
|
<p>{s.suggestion}</p>
|
|
{devMode && (
|
|
<span className="bottom-0 absolute text-zinc-700 dark:text-zinc-400 text-sm leading-10 h-10 right-2">
|
|
{s.relevance}
|
|
</span>
|
|
)}
|
|
</PlainText>
|
|
);
|
|
} else if (s.type === "inpage-link") {
|
|
return (
|
|
<LinkSuggestion
|
|
key={i}
|
|
query={s.suggestion}
|
|
selected={i == selected}
|
|
inPage={true}
|
|
>
|
|
{s.prompt && (
|
|
<span className="text-zinc-700 dark:text-zinc-400">{s.prompt}</span>
|
|
)}
|
|
{s.suggestion}
|
|
{devMode && (
|
|
<span className="absolute text-zinc-700 dark:text-zinc-400 text-sm leading-10 h-10 right-2">
|
|
{s.relevance}
|
|
</span>
|
|
)}
|
|
</LinkSuggestion>
|
|
);
|
|
}
|
|
})}
|
|
</SuggestionBox>
|
|
);
|
|
}
|