mirror of
https://github.com/hwchase17/langchain.git
synced 2025-08-08 04:25:46 +00:00
nomic: implement local embeddings with the inference_mode parameter (#21934)
## Description This PR implements local and dynamic mode in the Nomic Embed integration using the inference_mode and device parameters. They work as documented [here](https://docs.nomic.ai/reference/python-api/embeddings#local-inference). <!-- If no one reviews your PR within a few days, please @-mention one of baskaryan, efriis, eyurtsev, hwchase17. --> --------- Co-authored-by: Erick Friis <erickfriis@gmail.com>
This commit is contained in:
parent
0e72ed39a0
commit
25d1c1c9bb
@ -1,9 +1,9 @@
|
|||||||
import os
|
import os
|
||||||
from typing import List, Optional
|
from typing import List, Literal, Optional, overload
|
||||||
|
|
||||||
import nomic # type: ignore
|
import nomic # type: ignore[import]
|
||||||
from langchain_core.embeddings import Embeddings
|
from langchain_core.embeddings import Embeddings
|
||||||
from nomic import embed # type: ignore
|
from nomic import embed
|
||||||
|
|
||||||
|
|
||||||
class NomicEmbeddings(Embeddings):
|
class NomicEmbeddings(Embeddings):
|
||||||
@ -17,12 +17,46 @@ class NomicEmbeddings(Embeddings):
|
|||||||
model = NomicEmbeddings()
|
model = NomicEmbeddings()
|
||||||
"""
|
"""
|
||||||
|
|
||||||
|
@overload
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
model: str,
|
||||||
|
dimensionality: Optional[int] = ...,
|
||||||
|
inference_mode: Literal["remote"] = ...,
|
||||||
|
):
|
||||||
|
...
|
||||||
|
|
||||||
|
@overload
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
model: str,
|
||||||
|
dimensionality: Optional[int] = ...,
|
||||||
|
inference_mode: Literal["local", "dynamic"],
|
||||||
|
device: Optional[str] = ...,
|
||||||
|
):
|
||||||
|
...
|
||||||
|
|
||||||
|
@overload
|
||||||
|
def __init__(
|
||||||
|
self,
|
||||||
|
*,
|
||||||
|
model: str,
|
||||||
|
dimensionality: Optional[int] = ...,
|
||||||
|
inference_mode: str,
|
||||||
|
device: Optional[str] = ...,
|
||||||
|
):
|
||||||
|
...
|
||||||
|
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
*,
|
*,
|
||||||
model: str,
|
model: str,
|
||||||
nomic_api_key: Optional[str] = None,
|
nomic_api_key: Optional[str] = None,
|
||||||
dimensionality: Optional[int] = None,
|
dimensionality: Optional[int] = None,
|
||||||
|
inference_mode: str = "remote",
|
||||||
|
device: Optional[str] = None,
|
||||||
):
|
):
|
||||||
"""Initialize NomicEmbeddings model.
|
"""Initialize NomicEmbeddings model.
|
||||||
|
|
||||||
@ -30,12 +64,22 @@ class NomicEmbeddings(Embeddings):
|
|||||||
model: model name
|
model: model name
|
||||||
nomic_api_key: optionally, set the Nomic API key. Uses the NOMIC_API_KEY
|
nomic_api_key: optionally, set the Nomic API key. Uses the NOMIC_API_KEY
|
||||||
environment variable by default.
|
environment variable by default.
|
||||||
|
dimensionality: The embedding dimension, for use with Matryoshka-capable
|
||||||
|
models. Defaults to full-size.
|
||||||
|
inference_mode: How to generate embeddings. One of `remote`, `local`
|
||||||
|
(Embed4All), or `dynamic` (automatic). Defaults to `remote`.
|
||||||
|
device: The device to use for local embeddings. Choices include
|
||||||
|
`cpu`, `gpu`, `nvidia`, `amd`, or a specific device name. See
|
||||||
|
the docstring for `GPT4All.__init__` for more info. Typically
|
||||||
|
defaults to CPU. Do not use on macOS.
|
||||||
"""
|
"""
|
||||||
_api_key = nomic_api_key or os.environ.get("NOMIC_API_KEY")
|
_api_key = nomic_api_key or os.environ.get("NOMIC_API_KEY")
|
||||||
if _api_key:
|
if _api_key:
|
||||||
nomic.login(_api_key)
|
nomic.login(_api_key)
|
||||||
self.model = model
|
self.model = model
|
||||||
self.dimensionality = dimensionality
|
self.dimensionality = dimensionality
|
||||||
|
self.inference_mode = inference_mode
|
||||||
|
self.device = device
|
||||||
|
|
||||||
def embed(self, texts: List[str], *, task_type: str) -> List[List[float]]:
|
def embed(self, texts: List[str], *, task_type: str) -> List[List[float]]:
|
||||||
"""Embed texts.
|
"""Embed texts.
|
||||||
@ -51,6 +95,8 @@ class NomicEmbeddings(Embeddings):
|
|||||||
model=self.model,
|
model=self.model,
|
||||||
task_type=task_type,
|
task_type=task_type,
|
||||||
dimensionality=self.dimensionality,
|
dimensionality=self.dimensionality,
|
||||||
|
inference_mode=self.inference_mode,
|
||||||
|
device=self.device,
|
||||||
)
|
)
|
||||||
return output["embeddings"]
|
return output["embeddings"]
|
||||||
|
|
||||||
|
10
libs/partners/nomic/poetry.lock
generated
10
libs/partners/nomic/poetry.lock
generated
@ -1,4 +1,4 @@
|
|||||||
# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
|
# This file is automatically @generated by Poetry 1.8.2 and should not be changed by hand.
|
||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "annotated-types"
|
name = "annotated-types"
|
||||||
@ -430,12 +430,12 @@ files = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "nomic"
|
name = "nomic"
|
||||||
version = "3.0.28"
|
version = "3.0.29"
|
||||||
description = "The official Nomic python client."
|
description = "The official Nomic python client."
|
||||||
optional = false
|
optional = false
|
||||||
python-versions = "*"
|
python-versions = "*"
|
||||||
files = [
|
files = [
|
||||||
{file = "nomic-3.0.28.tar.gz", hash = "sha256:5712c58f0dcc8c446c1a30587490fb959a86398ad7a8e41bd68bd5f14096386b"},
|
{file = "nomic-3.0.29.tar.gz", hash = "sha256:1f92912e93932797f59dec8a33493f4cca92831aa86110590543f628098be1ef"},
|
||||||
]
|
]
|
||||||
|
|
||||||
[package.dependencies]
|
[package.dependencies]
|
||||||
@ -598,8 +598,8 @@ files = [
|
|||||||
[package.dependencies]
|
[package.dependencies]
|
||||||
numpy = [
|
numpy = [
|
||||||
{version = ">=1.20.3", markers = "python_version < \"3.10\""},
|
{version = ">=1.20.3", markers = "python_version < \"3.10\""},
|
||||||
{version = ">=1.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""},
|
|
||||||
{version = ">=1.23.2", markers = "python_version >= \"3.11\""},
|
{version = ">=1.23.2", markers = "python_version >= \"3.11\""},
|
||||||
|
{version = ">=1.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""},
|
||||||
]
|
]
|
||||||
python-dateutil = ">=2.8.2"
|
python-dateutil = ">=2.8.2"
|
||||||
pytz = ">=2020.1"
|
pytz = ">=2020.1"
|
||||||
@ -1309,4 +1309,4 @@ dev = ["black (>=19.3b0)", "pytest (>=4.6.2)"]
|
|||||||
[metadata]
|
[metadata]
|
||||||
lock-version = "2.0"
|
lock-version = "2.0"
|
||||||
python-versions = ">=3.8.1,<4.0"
|
python-versions = ">=3.8.1,<4.0"
|
||||||
content-hash = "0a9d4928b1fa0b379f6083bb054931c75c1f9cfe8e86bfd7d40ce484a993ac39"
|
content-hash = "369d2f7218d797a01a533380de9cce01037963f628dce10bc9927eac014edeeb"
|
||||||
|
@ -1,6 +1,6 @@
|
|||||||
[tool.poetry]
|
[tool.poetry]
|
||||||
name = "langchain-nomic"
|
name = "langchain-nomic"
|
||||||
version = "0.1.0"
|
version = "0.1.1"
|
||||||
description = "An integration package connecting Nomic and LangChain"
|
description = "An integration package connecting Nomic and LangChain"
|
||||||
authors = []
|
authors = []
|
||||||
readme = "README.md"
|
readme = "README.md"
|
||||||
@ -13,7 +13,7 @@ license = "MIT"
|
|||||||
[tool.poetry.dependencies]
|
[tool.poetry.dependencies]
|
||||||
python = ">=3.8.1,<4.0"
|
python = ">=3.8.1,<4.0"
|
||||||
langchain-core = ">=0.1.46,<0.3"
|
langchain-core = ">=0.1.46,<0.3"
|
||||||
nomic = "^3.0.12"
|
nomic = "^3.0.29"
|
||||||
|
|
||||||
[tool.poetry.group.test]
|
[tool.poetry.group.test]
|
||||||
optional = true
|
optional = true
|
||||||
|
Loading…
Reference in New Issue
Block a user