mincomk/risuai
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update transformers import and add config parameter to runTransformers function

Browse Source
This commit is contained in:
kwaroran
2024-01-05 23:54:01 +09:00
parent 5ae3f35df6
commit 56edb14e44
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
src/ts/process/embedding/transformers.ts
View File

@@ -1,4 +1,4 @@
import transformers, { AutoTokenizer, Pipeline, pipeline, type DataArray, type SummarizationOutput } from '@xenova/transformers'; import transformers, { AutoTokenizer, pipeline, type SummarizationOutput } from '@xenova/transformers';
transformers.env.localModelPath = "https://sv.risuai.xyz/transformers/" transformers.env.localModelPath = "https://sv.risuai.xyz/transformers/"
@@ -25,10 +25,10 @@ type TransformersBodyType = {
} }
export const runTransformers = async (baseText:string, model:string,bodyTemplate:TransformersBodyType) => { export const runTransformers = async (baseText:string, model:string,config:transformers.TextGenerationConfig = {}) => {
let text = baseText let text = baseText
let generator = await pipeline('text-generation', model); let generator = await pipeline('text-generation', model);
let output = await generator(text) as transformers.TextGenerationOutput let output = await generator(text, config) as transformers.TextGenerationOutput
const outputOne = output[0] const outputOne = output[0]
return outputOne return outputOne
} }