ref: adds multi threading to the AI/ML embeddings component (#2959)
* Use http client for requests and split texts naively * update models list * prints * multithread requests to aiml embeddings * remove comment * [autofix.ci] apply automated fixes * style(AIMLEmbeddingsImpl.py): improve code formatting and add type hinting for better code clarit --------- Co-authored-by: autofix-ci[bot] <114827586+autofix-ci[bot]@users.noreply.github.com> Co-authored-by: Gabriel Luiz Freitas Almeida <gabriel@langflow.org>
This commit is contained in:
Jordan Frazier
GitHub
parent
6a482f36c4
commit
55693c920d
7 changed files with 61 additions and 59 deletions
|
|
@ -1,4 +1,4 @@
|
|||
CHAT_MODELS = [
|
||||
AIML_CHAT_MODELS = [
|
||||
"zero-one-ai/Yi-34B-Chat",
|
||||
"allenai/OLMo-7B-Instruct",
|
||||
"allenai/OLMo-7B-Twin-2T",
|
||||
|
|
@ -17,6 +17,9 @@ CHAT_MODELS = [
|
|||
"codellama/CodeLlama-34b-Instruct-hf",
|
||||
"codellama/CodeLlama-70b-Instruct-hf",
|
||||
"codellama/CodeLlama-7b-Instruct-hf",
|
||||
"meta-llama/Meta-Llama-3.1-405B-Instruct-Turbo",
|
||||
"meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo",
|
||||
"meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo",
|
||||
"meta-llama/Llama-2-70b-chat-hf",
|
||||
"meta-llama/Llama-2-13b-chat-hf",
|
||||
"meta-llama/Llama-2-7b-chat-hf",
|
||||
|
|
@ -70,7 +73,7 @@ CHAT_MODELS = [
|
|||
"claude-3-haiku-20240307",
|
||||
]
|
||||
|
||||
EMBEDDING_MODELS = [
|
||||
AIML_EMBEDDING_MODELS = [
|
||||
"text-embedding-3-small",
|
||||
"text-embedding-3-large",
|
||||
"text-embedding-ada-002",
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
MODEL_NAMES = [
|
||||
OPENAI_MODEL_NAMES = [
|
||||
"gpt-4o-mini",
|
||||
"gpt-4o",
|
||||
"gpt-4-turbo",
|
||||
|
|
@ -7,3 +7,9 @@ MODEL_NAMES = [
|
|||
"gpt-3.5-turbo",
|
||||
"gpt-3.5-turbo-0125",
|
||||
]
|
||||
|
||||
OPENAI_EMBEDDING_MODEL_NAMES = [
|
||||
"text-embedding-3-small",
|
||||
"text-embedding-3-large",
|
||||
"text-embedding-ada-002",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
from langflow.base.embeddings.model import LCEmbeddingsModel
|
||||
from langflow.base.models.aiml_constants import EMBEDDING_MODELS
|
||||
from langflow.base.models.aiml_constants import AIML_EMBEDDING_MODELS
|
||||
from langflow.components.embeddings.util.AIMLEmbeddingsImpl import AIMLEmbeddingsImpl
|
||||
from langflow.field_typing import Embeddings
|
||||
from langflow.inputs.inputs import DropdownInput
|
||||
|
|
@ -16,7 +16,7 @@ class AIMLEmbeddingsComponent(LCEmbeddingsModel):
|
|||
DropdownInput(
|
||||
name="model_name",
|
||||
display_name="Model Name",
|
||||
options=EMBEDDING_MODELS,
|
||||
options=AIML_EMBEDDING_MODELS,
|
||||
required=True,
|
||||
),
|
||||
SecretStrInput(
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
from langchain_openai.embeddings.base import OpenAIEmbeddings
|
||||
|
||||
from langflow.base.embeddings.model import LCEmbeddingsModel
|
||||
from langflow.base.models.openai_constants import OPENAI_EMBEDDING_MODEL_NAMES
|
||||
from langflow.field_typing import Embeddings
|
||||
from langflow.io import BoolInput, DictInput, DropdownInput, FloatInput, IntInput, MessageTextInput, SecretStrInput
|
||||
|
||||
|
|
@ -33,11 +34,7 @@ class OpenAIEmbeddingsComponent(LCEmbeddingsModel):
|
|||
name="model",
|
||||
display_name="Model",
|
||||
advanced=False,
|
||||
options=[
|
||||
"text-embedding-3-small",
|
||||
"text-embedding-3-large",
|
||||
"text-embedding-ada-002",
|
||||
],
|
||||
options=OPENAI_EMBEDDING_MODEL_NAMES,
|
||||
value="text-embedding-3-small",
|
||||
),
|
||||
DictInput(name="model_kwargs", display_name="Model Kwargs", advanced=True),
|
||||
|
|
|
|||
|
|
@ -1,12 +1,13 @@
|
|||
import concurrent.futures
|
||||
import json
|
||||
from typing import List
|
||||
|
||||
import httpx
|
||||
from langflow.field_typing import Embeddings
|
||||
from langchain_core.runnables.config import run_in_executor
|
||||
from langchain_core.pydantic_v1 import BaseModel, SecretStr
|
||||
from loguru import logger
|
||||
|
||||
from langflow.field_typing import Embeddings
|
||||
|
||||
|
||||
class AIMLEmbeddingsImpl(BaseModel, Embeddings):
|
||||
embeddings_completion_url: str = "https://api.aimlapi.com/v1/embeddings"
|
||||
|
|
@ -15,56 +16,47 @@ class AIMLEmbeddingsImpl(BaseModel, Embeddings):
|
|||
model: str
|
||||
|
||||
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||||
result_vectors = []
|
||||
for text in texts:
|
||||
vector = self.embed_query(text)
|
||||
result_vectors.append(vector)
|
||||
|
||||
return result_vectors
|
||||
|
||||
def embed_query(self, text: str) -> List[float]:
|
||||
embeddings = [None] * len(texts)
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {self.api_key.get_secret_value()}",
|
||||
}
|
||||
|
||||
with httpx.Client() as client:
|
||||
with concurrent.futures.ThreadPoolExecutor() as executor:
|
||||
futures = []
|
||||
for i, text in enumerate(texts):
|
||||
futures.append((i, executor.submit(self._embed_text, client, headers, text)))
|
||||
|
||||
for index, future in futures:
|
||||
try:
|
||||
result_data = future.result()
|
||||
assert len(result_data["data"]) == 1, "Expected one embedding"
|
||||
embeddings[index] = result_data["data"][0]["embedding"]
|
||||
except (
|
||||
httpx.HTTPStatusError,
|
||||
httpx.RequestError,
|
||||
json.JSONDecodeError,
|
||||
KeyError,
|
||||
) as e:
|
||||
logger.error(f"Error occurred: {e}")
|
||||
raise
|
||||
|
||||
return embeddings # type: ignore
|
||||
|
||||
def _embed_text(self, client: httpx.Client, headers: dict, text: str) -> dict:
|
||||
payload = {
|
||||
"model": self.model,
|
||||
"input": text,
|
||||
}
|
||||
vector = []
|
||||
try:
|
||||
response = httpx.post(
|
||||
self.embeddings_completion_url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
)
|
||||
try:
|
||||
response.raise_for_status()
|
||||
result_data = response.json()
|
||||
vector = result_data["data"][0]["embedding"]
|
||||
except httpx.HTTPStatusError as http_err:
|
||||
logger.error(f"HTTP error occurred: {http_err}")
|
||||
raise http_err
|
||||
except httpx.RequestError as req_err:
|
||||
logger.error(f"Request error occurred: {req_err}")
|
||||
raise req_err
|
||||
except json.JSONDecodeError:
|
||||
logger.warning(f"Failed to decode JSON, response text: {response.text}")
|
||||
except KeyError as key_err:
|
||||
logger.warning(f"Key error: {key_err}, response content: {result_data}")
|
||||
raise key_err
|
||||
except httpx.TimeoutException:
|
||||
logger.error("Request timed out.")
|
||||
raise
|
||||
except Exception as exc:
|
||||
logger.error(f"Error: {exc}")
|
||||
raise
|
||||
response = client.post(
|
||||
self.embeddings_completion_url,
|
||||
headers=headers,
|
||||
json=payload,
|
||||
)
|
||||
response.raise_for_status()
|
||||
result_data = response.json()
|
||||
return result_data
|
||||
|
||||
return vector
|
||||
|
||||
async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
|
||||
return await run_in_executor(None, self.embed_documents, texts)
|
||||
|
||||
async def aembed_query(self, text: str) -> List[float]:
|
||||
return await run_in_executor(None, self.embed_query, text)
|
||||
def embed_query(self, text: str) -> List[float]:
|
||||
return self.embed_documents([text])[0]
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import json
|
||||
import httpx
|
||||
from langflow.base.models.aiml_constants import CHAT_MODELS
|
||||
from langflow.base.models.aiml_constants import AIML_CHAT_MODELS
|
||||
from langflow.custom.custom_component.component import Component
|
||||
|
||||
from langflow.inputs.inputs import FloatInput, IntInput, MessageInput, SecretStrInput
|
||||
|
|
@ -29,7 +29,7 @@ class AIMLModelComponent(Component):
|
|||
DropdownInput(
|
||||
name="model_name",
|
||||
display_name="Model Name",
|
||||
options=CHAT_MODELS,
|
||||
options=AIML_CHAT_MODELS,
|
||||
required=True,
|
||||
),
|
||||
SecretStrInput(
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ from langchain_openai import ChatOpenAI
|
|||
from pydantic.v1 import SecretStr
|
||||
|
||||
from langflow.base.models.model import LCModelComponent
|
||||
from langflow.base.models.openai_constants import MODEL_NAMES
|
||||
from langflow.base.models.openai_constants import OPENAI_MODEL_NAMES
|
||||
from langflow.field_typing import LanguageModel
|
||||
from langflow.inputs import (
|
||||
BoolInput,
|
||||
|
|
@ -46,7 +46,11 @@ class OpenAIModelComponent(LCModelComponent):
|
|||
info="The schema for the Output of the model. You must pass the word JSON in the prompt. If left blank, JSON mode will be disabled.",
|
||||
),
|
||||
DropdownInput(
|
||||
name="model_name", display_name="Model Name", advanced=False, options=MODEL_NAMES, value=MODEL_NAMES[0]
|
||||
name="model_name",
|
||||
display_name="Model Name",
|
||||
advanced=False,
|
||||
options=OPENAI_MODEL_NAMES,
|
||||
value=OPENAI_MODEL_NAMES[0],
|
||||
),
|
||||
StrInput(
|
||||
name="openai_api_base",
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue