Add new OpenAI embedding models for selection

This commit is contained in:
Cohee 2024-03-12 21:41:30 +02:00
parent 44a7dd3d74
commit 5efdecc825
3 changed files with 65 additions and 13 deletions

View File

@ -13,6 +13,7 @@ const settings = {
source: 'transformers',
include_wi: false,
togetherai_model: 'togethercomputer/m2-bert-80M-32k-retrieval',
openai_model: 'text-embedding-ada-002',
// For chats
enabled_chats: false,
@ -429,6 +430,24 @@ async function getSavedHashes(collectionId) {
return hashes;
}
function getVectorHeaders() {
const headers = getRequestHeaders();
switch (settings.source) {
case 'extras':
addExtrasHeaders(headers);
break;
case 'togetherai':
addTogetherAiHeaders(headers);
break;
case 'openai':
addOpenAiHeaders(headers);
break;
default:
break;
}
return headers;
}
/**
* Add headers for the Extras API source.
* @param {object} headers Headers object
@ -442,7 +461,7 @@ function addExtrasHeaders(headers) {
}
/**
* Add headers for the Extras API source.
* Add headers for the TogetherAI API source.
* @param {object} headers Headers object
*/
function addTogetherAiHeaders(headers) {
@ -451,6 +470,16 @@ function addTogetherAiHeaders(headers) {
});
}
/**
* Add headers for the OpenAI API source.
* @param {object} headers Header object
*/
function addOpenAiHeaders(headers) {
Object.assign(headers, {
'X-OpenAI-Model': extension_settings.vectors.openai_model,
});
}
/**
* Inserts vector items into a collection
* @param {string} collectionId - The collection to insert into
@ -470,12 +499,7 @@ async function insertVectorItems(collectionId, items) {
throw new Error('Vectors: Embeddings module missing', { cause: 'extras_module_missing' });
}
const headers = getRequestHeaders();
if (settings.source === 'extras') {
addExtrasHeaders(headers);
} else if (settings.source === 'togetherai') {
addTogetherAiHeaders(headers);
}
const headers = getVectorHeaders();
const response = await fetch('/api/vector/insert', {
method: 'POST',
@ -521,12 +545,7 @@ async function deleteVectorItems(collectionId, hashes) {
* @returns {Promise<{ hashes: number[], metadata: object[]}>} - Hashes of the results
*/
async function queryCollection(collectionId, searchText, topK) {
const headers = getRequestHeaders();
if (settings.source === 'extras') {
addExtrasHeaders(headers);
} else if (settings.source === 'togetherai') {
addTogetherAiHeaders(headers);
}
const headers = getVectorHeaders();
const response = await fetch('/api/vector/query', {
method: 'POST',
@ -581,6 +600,7 @@ function toggleSettings() {
$('#vectors_files_settings').toggle(!!settings.enabled_files);
$('#vectors_chats_settings').toggle(!!settings.enabled_chats);
$('#together_vectorsModel').toggle(settings.source === 'togetherai');
$('#openai_vectorsModel').toggle(settings.source === 'openai');
$('#nomicai_apiKey').toggle(settings.source === 'nomicai');
}
@ -645,6 +665,7 @@ jQuery(async () => {
saveSettingsDebounced();
toggleSettings();
});
$('#vectors_modelWarning').hide();
$('#vectors_enabled_files').prop('checked', settings.enabled_files).on('input', () => {
settings.enabled_files = $('#vectors_enabled_files').prop('checked');
Object.assign(extension_settings.vectors, settings);
@ -665,10 +686,17 @@ jQuery(async () => {
saveSettingsDebounced();
});
$('#vectors_togetherai_model').val(settings.togetherai_model).on('change', () => {
$('#vectors_modelWarning').show();
settings.togetherai_model = String($('#vectors_togetherai_model').val());
Object.assign(extension_settings.vectors, settings);
saveSettingsDebounced();
});
$('#vectors_openai_model').val(settings.openai_model).on('change', () => {
$('#vectors_modelWarning').show();
settings.openai_model = String($('#vectors_openai_model').val());
Object.assign(extension_settings.vectors, settings);
saveSettingsDebounced();
});
$('#vectors_template').val(settings.template).on('input', () => {
settings.template = String($('#vectors_template').val());
Object.assign(extension_settings.vectors, settings);

View File

@ -19,6 +19,16 @@
<option value="nomicai">NomicAI</option>
</select>
</div>
<div class="flex-container flexFlowColumn" id="openai_vectorsModel">
<label for="vectors_openai_model">
Vectorization Model
</label>
<select id="vectors_openai_model" class="text_pole">
<option value="text-embedding-ada-002">text-embedding-ada-002</option>
<option value="text-embedding-3-small">text-embedding-3-small</option>
<option value="text-embedding-3-large">text-embedding-3-large</option>
</select>
</div>
<div class="flex-container flexFlowColumn" id="together_vectorsModel">
<label for="vectors_togetherai_model">
Vectorization Model
@ -34,6 +44,14 @@
<option value="bert-base-uncased">Bert Base Uncased</option>
</select>
</div>
<small id="vectors_modelWarning">
<i class="fa-solid fa-exclamation-triangle"></i>
<span data-i18n="Vectors Model Warning">
It is recommended to purge vectors when changing the model mid-chat. Otherwise, it will lead to sub-par results.
</span>
</small>
<div class="flex-container flexFlowColumn" id="nomicai_apiKey">
<label for="api_key_nomicai">
<span>NomicAI API Key</span>

View File

@ -175,6 +175,12 @@ function getSourceSettings(source, request) {
if (source === 'togetherai') {
let model = String(request.headers['x-togetherai-model']);
return {
model: model,
};
} else if (source === 'openai') {
let model = String(request.headers['x-openai-model']);
return {
model: model,
};