Merge pull request #2785 from cheahjs/feat/openai-embeddings-batch

feat: add RAG_EMBEDDING_OPENAI_BATCH_SIZE to batch multiple embeddings
This commit is contained in:
Timothy Jaeryang Baek 2024-06-03 13:50:14 -07:00 committed by GitHub
commit 768941bded
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
39 changed files with 112 additions and 19 deletions

View File

@ -78,6 +78,7 @@ from utils.misc import (
from utils.utils import get_current_user, get_admin_user
from config import (
AppConfig,
ENV,
SRC_LOG_LEVELS,
UPLOAD_DIR,
@ -114,7 +115,7 @@ from config import (
SERPER_API_KEY,
RAG_WEB_SEARCH_RESULT_COUNT,
RAG_WEB_SEARCH_CONCURRENT_REQUESTS,
AppConfig,
RAG_EMBEDDING_OPENAI_BATCH_SIZE,
)
from constants import ERROR_MESSAGES
@ -139,6 +140,7 @@ app.state.config.CHUNK_OVERLAP = CHUNK_OVERLAP
app.state.config.RAG_EMBEDDING_ENGINE = RAG_EMBEDDING_ENGINE
app.state.config.RAG_EMBEDDING_MODEL = RAG_EMBEDDING_MODEL
app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = RAG_EMBEDDING_OPENAI_BATCH_SIZE
app.state.config.RAG_RERANKING_MODEL = RAG_RERANKING_MODEL
app.state.config.RAG_TEMPLATE = RAG_TEMPLATE
@ -212,6 +214,7 @@ app.state.EMBEDDING_FUNCTION = get_embedding_function(
app.state.sentence_transformer_ef,
app.state.config.OPENAI_API_KEY,
app.state.config.OPENAI_API_BASE_URL,
app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
)
origins = ["*"]
@ -248,6 +251,7 @@ async def get_status():
"embedding_engine": app.state.config.RAG_EMBEDDING_ENGINE,
"embedding_model": app.state.config.RAG_EMBEDDING_MODEL,
"reranking_model": app.state.config.RAG_RERANKING_MODEL,
"openai_batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
}
@ -260,6 +264,7 @@ async def get_embedding_config(user=Depends(get_admin_user)):
"openai_config": {
"url": app.state.config.OPENAI_API_BASE_URL,
"key": app.state.config.OPENAI_API_KEY,
"batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
},
}
@ -275,6 +280,7 @@ async def get_reraanking_config(user=Depends(get_admin_user)):
class OpenAIConfigForm(BaseModel):
url: str
key: str
batch_size: Optional[int] = None
class EmbeddingModelUpdateForm(BaseModel):
@ -295,9 +301,14 @@ async def update_embedding_config(
app.state.config.RAG_EMBEDDING_MODEL = form_data.embedding_model
if app.state.config.RAG_EMBEDDING_ENGINE in ["ollama", "openai"]:
if form_data.openai_config != None:
if form_data.openai_config is not None:
app.state.config.OPENAI_API_BASE_URL = form_data.openai_config.url
app.state.config.OPENAI_API_KEY = form_data.openai_config.key
app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE = (
form_data.openai_config.batch_size
if form_data.openai_config.batch_size
else 1
)
update_embedding_model(app.state.config.RAG_EMBEDDING_MODEL)
@ -307,6 +318,7 @@ async def update_embedding_config(
app.state.sentence_transformer_ef,
app.state.config.OPENAI_API_KEY,
app.state.config.OPENAI_API_BASE_URL,
app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
)
return {
@ -316,6 +328,7 @@ async def update_embedding_config(
"openai_config": {
"url": app.state.config.OPENAI_API_BASE_URL,
"key": app.state.config.OPENAI_API_KEY,
"batch_size": app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
},
}
except Exception as e:
@ -881,6 +894,7 @@ def store_docs_in_vector_db(docs, collection_name, overwrite: bool = False) -> b
app.state.sentence_transformer_ef,
app.state.config.OPENAI_API_KEY,
app.state.config.OPENAI_API_BASE_URL,
app.state.config.RAG_EMBEDDING_OPENAI_BATCH_SIZE,
)
embedding_texts = list(map(lambda x: x.replace("\n", " "), texts))

View File

@ -2,7 +2,7 @@ import os
import logging
import requests
from typing import List
from typing import List, Union
from apps.ollama.main import (
generate_ollama_embeddings,
@ -21,17 +21,7 @@ from langchain.retrievers import (
from typing import Optional
from config import (
SRC_LOG_LEVELS,
CHROMA_CLIENT,
SEARXNG_QUERY_URL,
GOOGLE_PSE_API_KEY,
GOOGLE_PSE_ENGINE_ID,
BRAVE_SEARCH_API_KEY,
SERPSTACK_API_KEY,
SERPSTACK_HTTPS,
SERPER_API_KEY,
)
from config import SRC_LOG_LEVELS, CHROMA_CLIENT
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["RAG"])
@ -209,6 +199,7 @@ def get_embedding_function(
embedding_function,
openai_key,
openai_url,
batch_size,
):
if embedding_engine == "":
return lambda query: embedding_function.encode(query).tolist()
@ -232,7 +223,13 @@ def get_embedding_function(
def generate_multiple(query, f):
if isinstance(query, list):
return [f(q) for q in query]
if embedding_engine == "openai":
embeddings = []
for i in range(0, len(query), batch_size):
embeddings.extend(f(query[i : i + batch_size]))
return embeddings
else:
return [f(q) for q in query]
else:
return f(query)
@ -413,8 +410,22 @@ def get_model_path(model: str, update_model: bool = False):
def generate_openai_embeddings(
model: str, text: str, key: str, url: str = "https://api.openai.com/v1"
model: str,
text: Union[str, list[str]],
key: str,
url: str = "https://api.openai.com/v1",
):
if isinstance(text, list):
embeddings = generate_openai_batch_embeddings(model, text, key, url)
else:
embeddings = generate_openai_batch_embeddings(model, [text], key, url)
return embeddings[0] if isinstance(text, str) else embeddings
def generate_openai_batch_embeddings(
model: str, texts: list[str], key: str, url: str = "https://api.openai.com/v1"
) -> Optional[list[list[float]]]:
try:
r = requests.post(
f"{url}/embeddings",
@ -422,12 +433,12 @@ def generate_openai_embeddings(
"Content-Type": "application/json",
"Authorization": f"Bearer {key}",
},
json={"input": text, "model": model},
json={"input": texts, "model": model},
)
r.raise_for_status()
data = r.json()
if "data" in data:
return data["data"][0]["embedding"]
return [elem["embedding"] for elem in data["data"]]
else:
raise "Something went wrong :/"
except Exception as e:

View File

@ -683,6 +683,12 @@ RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE = (
os.environ.get("RAG_EMBEDDING_MODEL_TRUST_REMOTE_CODE", "").lower() == "true"
)
RAG_EMBEDDING_OPENAI_BATCH_SIZE = PersistentConfig(
"RAG_EMBEDDING_OPENAI_BATCH_SIZE",
"rag.embedding_openai_batch_size",
os.environ.get("RAG_EMBEDDING_OPENAI_BATCH_SIZE", 1),
)
RAG_RERANKING_MODEL = PersistentConfig(
"RAG_RERANKING_MODEL",
"rag.reranking_model",

View File

@ -415,6 +415,7 @@ export const getEmbeddingConfig = async (token: string) => {
type OpenAIConfigForm = {
key: string;
url: string;
batch_size: number;
};
type EmbeddingModelUpdateForm = {

View File

@ -31,6 +31,7 @@
let OpenAIKey = '';
let OpenAIUrl = '';
let OpenAIBatchSize = 1;
let querySettings = {
template: '',
@ -92,7 +93,8 @@
? {
openai_config: {
key: OpenAIKey,
url: OpenAIUrl
url: OpenAIUrl,
batch_size: OpenAIBatchSize
}
}
: {})
@ -159,6 +161,7 @@
OpenAIKey = embeddingConfig.openai_config.key;
OpenAIUrl = embeddingConfig.openai_config.url;
OpenAIBatchSize = embeddingConfig.openai_config.batch_size ?? 1;
}
};
@ -282,6 +285,30 @@
required
/>
</div>
<div class="flex mt-0.5 space-x-2">
<div class=" self-center text-xs font-medium">{$i18n.t('Embedding Batch Size')}</div>
<div class=" flex-1">
<input
id="steps-range"
type="range"
min="1"
max="2048"
step="1"
bind:value={OpenAIBatchSize}
class="w-full h-2 rounded-lg appearance-none cursor-pointer dark:bg-gray-700"
/>
</div>
<div class="">
<input
bind:value={OpenAIBatchSize}
type="number"
class=" bg-transparent text-center w-14"
min="-2"
max="16000"
step="1"
/>
</div>
</div>
{/if}
<div class=" flex w-full justify-between">

View File

@ -170,6 +170,7 @@
"Edit Doc": "تعديل الملف",
"Edit User": "تعديل المستخدم",
"Email": "البريد",
"Embedding Batch Size": "",
"Embedding Model": "نموذج التضمين",
"Embedding Model Engine": "تضمين محرك النموذج",
"Embedding model set to \"{{embedding_model}}\"": "تم تعيين نموذج التضمين على \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Редактиране на документ",
"Edit User": "Редактиране на потребител",
"Email": "Имейл",
"Embedding Batch Size": "",
"Embedding Model": "Модел за вграждане",
"Embedding Model Engine": "Модел за вграждане",
"Embedding model set to \"{{embedding_model}}\"": "Модел за вграждане е настроен на \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "ডকুমেন্ট এডিট করুন",
"Edit User": "ইউজার এডিট করুন",
"Email": "ইমেইল",
"Embedding Batch Size": "",
"Embedding Model": "ইমেজ ইমেবডিং মডেল",
"Embedding Model Engine": "ইমেজ ইমেবডিং মডেল ইঞ্জিন",
"Embedding model set to \"{{embedding_model}}\"": "ইমেজ ইমেবডিং মডেল সেট করা হয়েছে - \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Edita Document",
"Edit User": "Edita Usuari",
"Email": "Correu electrònic",
"Embedding Batch Size": "",
"Embedding Model": "Model d'embutiment",
"Embedding Model Engine": "Motor de model d'embutiment",
"Embedding model set to \"{{embedding_model}}\"": "Model d'embutiment configurat a \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "I-edit ang dokumento",
"Edit User": "I-edit ang tiggamit",
"Email": "E-mail",
"Embedding Batch Size": "",
"Embedding Model": "",
"Embedding Model Engine": "",
"Embedding model set to \"{{embedding_model}}\"": "",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Dokument bearbeiten",
"Edit User": "Benutzer bearbeiten",
"Email": "E-Mail",
"Embedding Batch Size": "",
"Embedding Model": "Embedding-Modell",
"Embedding Model Engine": "Embedding-Modell-Engine",
"Embedding model set to \"{{embedding_model}}\"": "Embedding-Modell auf \"{{embedding_model}}\" gesetzt",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Edit Doge",
"Edit User": "Edit Wowser",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "",
"Embedding Model Engine": "",
"Embedding model set to \"{{embedding_model}}\"": "",

View File

@ -170,6 +170,7 @@
"Edit Doc": "",
"Edit User": "",
"Email": "",
"Embedding Batch Size": "",
"Embedding Model": "",
"Embedding Model Engine": "",
"Embedding model set to \"{{embedding_model}}\"": "",

View File

@ -170,6 +170,7 @@
"Edit Doc": "",
"Edit User": "",
"Email": "",
"Embedding Batch Size": "",
"Embedding Model": "",
"Embedding Model Engine": "",
"Embedding model set to \"{{embedding_model}}\"": "",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Editar Documento",
"Edit User": "Editar Usuario",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Modelo de Embedding",
"Embedding Model Engine": "Motor de Modelo de Embedding",
"Embedding model set to \"{{embedding_model}}\"": "Modelo de Embedding configurado a \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "ویرایش سند",
"Edit User": "ویرایش کاربر",
"Email": "ایمیل",
"Embedding Batch Size": "",
"Embedding Model": "مدل پیدائش",
"Embedding Model Engine": "محرک مدل پیدائش",
"Embedding model set to \"{{embedding_model}}\"": "مدل پیدائش را به \"{{embedding_model}}\" تنظیم کنید",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Muokkaa asiakirjaa",
"Edit User": "Muokkaa käyttäjää",
"Email": "Sähköposti",
"Embedding Batch Size": "",
"Embedding Model": "Upotusmalli",
"Embedding Model Engine": "Upotusmallin moottori",
"Embedding model set to \"{{embedding_model}}\"": "\"{{embedding_model}}\" valittu upotusmalliksi",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Éditer le document",
"Edit User": "Éditer l'utilisateur",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Modèle d'embedding",
"Embedding Model Engine": "Moteur du modèle d'embedding",
"Embedding model set to \"{{embedding_model}}\"": "Modèle d'embedding défini sur \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Éditer le document",
"Edit User": "Éditer l'utilisateur",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Modèle pour l'Embedding",
"Embedding Model Engine": "Moteur du Modèle d'Embedding",
"Embedding model set to \"{{embedding_model}}\"": "Modèle d'embedding défini sur \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "ערוך מסמך",
"Edit User": "ערוך משתמש",
"Email": "דוא\"ל",
"Embedding Batch Size": "",
"Embedding Model": "מודל הטמעה",
"Embedding Model Engine": "מנוע מודל הטמעה",
"Embedding model set to \"{{embedding_model}}\"": "מודל ההטמעה הוגדר ל-\"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "दस्तावेज़ संपादित करें",
"Edit User": "यूजर को संपादित करो",
"Email": "ईमेल",
"Embedding Batch Size": "",
"Embedding Model": "मॉडेल अनुकूलन",
"Embedding Model Engine": "एंबेडिंग मॉडल इंजन",
"Embedding model set to \"{{embedding_model}}\"": "एम्बेडिंग मॉडल को \"{{embedding_model}}\" पर सेट किया गया",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Uredi dokument",
"Edit User": "Uredi korisnika",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Embedding model",
"Embedding Model Engine": "Embedding model pogon",
"Embedding model set to \"{{embedding_model}}\"": "Embedding model postavljen na \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Modifica documento",
"Edit User": "Modifica utente",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Modello di embedding",
"Embedding Model Engine": "Motore del modello di embedding",
"Embedding model set to \"{{embedding_model}}\"": "Modello di embedding impostato su \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "ドキュメントを編集",
"Edit User": "ユーザーを編集",
"Email": "メールアドレス",
"Embedding Batch Size": "",
"Embedding Model": "埋め込みモデル",
"Embedding Model Engine": "埋め込みモデルエンジン",
"Embedding model set to \"{{embedding_model}}\"": "埋め込みモデルを\"{{embedding_model}}\"に設定しました",

View File

@ -170,6 +170,7 @@
"Edit Doc": "დოკუმენტის ედიტირება",
"Edit User": "მომხმარებლის ედიტირება",
"Email": "ელ-ფოსტა",
"Embedding Batch Size": "",
"Embedding Model": "ჩასმის ძირითადი პროგრამა",
"Embedding Model Engine": "ჩასმის ძირითადი პროგრამა",
"Embedding model set to \"{{embedding_model}}\"": "ჩასმის ძირითადი პროგრამა ჩართულია \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "문서 편집",
"Edit User": "사용자 편집",
"Email": "이메일",
"Embedding Batch Size": "",
"Embedding Model": "임베딩 모델",
"Embedding Model Engine": "임베딩 모델 엔진",
"Embedding model set to \"{{embedding_model}}\"": "임베딩 모델을 \"{{embedding_model}}\"로 설정됨",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Wijzig Doc",
"Edit User": "Wijzig Gebruiker",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Embedding Model",
"Embedding Model Engine": "Embedding Model Engine",
"Embedding model set to \"{{embedding_model}}\"": "Embedding model ingesteld op \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "ਡਾਕੂਮੈਂਟ ਸੰਪਾਦਨ ਕਰੋ",
"Edit User": "ਉਪਭੋਗਤਾ ਸੰਪਾਦਨ ਕਰੋ",
"Email": "ਈਮੇਲ",
"Embedding Batch Size": "",
"Embedding Model": "ਐਮਬੈੱਡਿੰਗ ਮਾਡਲ",
"Embedding Model Engine": "ਐਮਬੈੱਡਿੰਗ ਮਾਡਲ ਇੰਜਣ",
"Embedding model set to \"{{embedding_model}}\"": "ਐਮਬੈੱਡਿੰਗ ਮਾਡਲ ਨੂੰ \"{{embedding_model}}\" 'ਤੇ ਸੈੱਟ ਕੀਤਾ ਗਿਆ",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Edytuj dokument",
"Edit User": "Edytuj użytkownika",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Model osadzania",
"Embedding Model Engine": "Silnik modelu osadzania",
"Embedding model set to \"{{embedding_model}}\"": "Model osadzania ustawiono na \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Editar Documento",
"Edit User": "Editar Usuário",
"Email": "E-mail",
"Embedding Batch Size": "",
"Embedding Model": "Modelo de Embedding",
"Embedding Model Engine": "Motor de Modelo de Embedding",
"Embedding model set to \"{{embedding_model}}\"": "Modelo de Embedding definido como \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Editar Documento",
"Edit User": "Editar Usuário",
"Email": "E-mail",
"Embedding Batch Size": "",
"Embedding Model": "Modelo de Embedding",
"Embedding Model Engine": "Motor de Modelo de Embedding",
"Embedding model set to \"{{embedding_model}}\"": "Modelo de Embedding definido como \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Редактировать документ",
"Edit User": "Редактировать пользователя",
"Email": "Электронная почта",
"Embedding Batch Size": "",
"Embedding Model": "Модель эмбеддинга",
"Embedding Model Engine": "Модель эмбеддинга",
"Embedding model set to \"{{embedding_model}}\"": "Эмбеддинг-модель установлена в \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Уреди документ",
"Edit User": "Уреди корисника",
"Email": "Е-пошта",
"Embedding Batch Size": "",
"Embedding Model": "Модел уградње",
"Embedding Model Engine": "Мотор модела уградње",
"Embedding model set to \"{{embedding_model}}\"": "Модел уградње подешен на \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Redigera dokument",
"Edit User": "Redigera användare",
"Email": "E-post",
"Embedding Batch Size": "",
"Embedding Model": "Embeddingsmodell",
"Embedding Model Engine": "Embeddingsmodellmotor",
"Embedding model set to \"{{embedding_model}}\"": "Embeddingsmodell inställd på \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Belgeyi Düzenle",
"Edit User": "Kullanıcıyı Düzenle",
"Email": "E-posta",
"Embedding Batch Size": "",
"Embedding Model": "Gömme Modeli",
"Embedding Model Engine": "Gömme Modeli Motoru",
"Embedding model set to \"{{embedding_model}}\"": "Gömme modeli \"{{embedding_model}}\" olarak ayarlandı",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Редагувати документ",
"Edit User": "Редагувати користувача",
"Email": "Електронна пошта",
"Embedding Batch Size": "",
"Embedding Model": "Модель вбудовування",
"Embedding Model Engine": "Двигун модели встраивания ",
"Embedding model set to \"{{embedding_model}}\"": "Встановлена модель вбудовування \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "Thay đổi tài liệu",
"Edit User": "Thay đổi thông tin người sử dụng",
"Email": "Email",
"Embedding Batch Size": "",
"Embedding Model": "Mô hình embedding",
"Embedding Model Engine": "Trình xử lý embedding",
"Embedding model set to \"{{embedding_model}}\"": "Mô hình embedding đã được thiết lập thành \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "编辑文档",
"Edit User": "编辑用户",
"Email": "邮箱地址",
"Embedding Batch Size": "",
"Embedding Model": "语义向量模型",
"Embedding Model Engine": "语义向量模型引擎",
"Embedding model set to \"{{embedding_model}}\"": "语义向量模型设置为 \"{{embedding_model}}\"",

View File

@ -170,6 +170,7 @@
"Edit Doc": "編輯文件",
"Edit User": "編輯使用者",
"Email": "電子郵件",
"Embedding Batch Size": "",
"Embedding Model": "嵌入模型",
"Embedding Model Engine": "嵌入模型引擎",
"Embedding model set to \"{{embedding_model}}\"": "嵌入模型已設定為 \"{{embedding_model}}\"",