rixty/SillyTavern
1
0
Fork 0
mirror of https://github.com/SillyTavern/SillyTavern.git synced 2025-06-05 21:59:27 +02:00
Code Issues Packages Projects Releases Wiki Activity

Split vector batches into 10s

Browse Source
This commit is contained in:
Cohee
2024-02-01 11:02:47 +02:00
parent 7618133d6d
commit 695b438c0d
1 changed files with 23 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
src/endpoints/vectors.js
View File

@@ -35,19 +35,31 @@ async function getVector(source, sourceSettings, text) {
* @returns {Promise<number[][]>} - The array of vectors for the texts * @returns {Promise<number[][]>} - The array of vectors for the texts
*/ */
async function getBatchVector(source, sourceSettings, texts) { async function getBatchVector(source, sourceSettings, texts) {
const batchSize = 10;
const batches = Array(Math.ceil(texts.length / batchSize)).fill(undefined).map((_, i) => texts.slice(i * batchSize, i * batchSize + batchSize));
let results = [];
for (let batch of batches) {
switch (source) { switch (source) {
case 'mistral': case 'mistral':
case 'openai': case 'openai':
return require('../openai-vectors').getOpenAIBatchVector(texts, source); results.push(...await require('../openai-vectors').getOpenAIBatchVector(batch, source));
break;
case 'transformers': case 'transformers':
return require('../embedding').getTransformersBatchVector(texts); results.push(...await require('../embedding').getTransformersBatchVector(batch));
break;
case 'extras': case 'extras':
return require('../extras-vectors').getExtrasBatchVector(texts, sourceSettings.extrasUrl, sourceSettings.extrasKey); results.push(...await require('../extras-vectors').getExtrasBatchVector(batch, sourceSettings.extrasUrl, sourceSettings.extrasKey));
break;
case 'palm': case 'palm':
return require('../makersuite-vectors').getMakerSuiteBatchVector(texts); results.push(...await require('../makersuite-vectors').getMakerSuiteBatchVector(batch));
break;
default:
throw new Error(`Unknown vector source ${source}`);
}
} }
throw new Error(`Unknown vector source ${source}`); return results;
} }
/** /**