diff --git a/libs/partners/nomic/langchain_nomic/embeddings.py b/libs/partners/nomic/langchain_nomic/embeddings.py index 55152c417f1..5c690276252 100644 --- a/libs/partners/nomic/langchain_nomic/embeddings.py +++ b/libs/partners/nomic/langchain_nomic/embeddings.py @@ -1,9 +1,9 @@ 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 nomic import embed # type: ignore +from nomic import embed class NomicEmbeddings(Embeddings): @@ -17,12 +17,46 @@ class NomicEmbeddings(Embeddings): 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__( self, *, model: str, nomic_api_key: Optional[str] = None, dimensionality: Optional[int] = None, + inference_mode: str = "remote", + device: Optional[str] = None, ): """Initialize NomicEmbeddings model. @@ -30,12 +64,22 @@ class NomicEmbeddings(Embeddings): model: model name nomic_api_key: optionally, set the Nomic API key. Uses the NOMIC_API_KEY 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") if _api_key: nomic.login(_api_key) self.model = model self.dimensionality = dimensionality + self.inference_mode = inference_mode + self.device = device def embed(self, texts: List[str], *, task_type: str) -> List[List[float]]: """Embed texts. @@ -51,6 +95,8 @@ class NomicEmbeddings(Embeddings): model=self.model, task_type=task_type, dimensionality=self.dimensionality, + inference_mode=self.inference_mode, + device=self.device, ) return output["embeddings"] diff --git a/libs/partners/nomic/poetry.lock b/libs/partners/nomic/poetry.lock index 2307d008c96..2e015a30a39 100644 --- a/libs/partners/nomic/poetry.lock +++ b/libs/partners/nomic/poetry.lock @@ -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]] name = "annotated-types" @@ -430,12 +430,12 @@ files = [ [[package]] name = "nomic" -version = "3.0.28" +version = "3.0.29" description = "The official Nomic python client." optional = false python-versions = "*" files = [ - {file = "nomic-3.0.28.tar.gz", hash = "sha256:5712c58f0dcc8c446c1a30587490fb959a86398ad7a8e41bd68bd5f14096386b"}, + {file = "nomic-3.0.29.tar.gz", hash = "sha256:1f92912e93932797f59dec8a33493f4cca92831aa86110590543f628098be1ef"}, ] [package.dependencies] @@ -598,8 +598,8 @@ files = [ [package.dependencies] numpy = [ {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.21.0", markers = "python_version >= \"3.10\" and python_version < \"3.11\""}, ] python-dateutil = ">=2.8.2" pytz = ">=2020.1" @@ -1309,4 +1309,4 @@ dev = ["black (>=19.3b0)", "pytest (>=4.6.2)"] [metadata] lock-version = "2.0" python-versions = ">=3.8.1,<4.0" -content-hash = "0a9d4928b1fa0b379f6083bb054931c75c1f9cfe8e86bfd7d40ce484a993ac39" +content-hash = "369d2f7218d797a01a533380de9cce01037963f628dce10bc9927eac014edeeb" diff --git a/libs/partners/nomic/pyproject.toml b/libs/partners/nomic/pyproject.toml index f2c0a339cfc..c3c822c830a 100644 --- a/libs/partners/nomic/pyproject.toml +++ b/libs/partners/nomic/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "langchain-nomic" -version = "0.1.0" +version = "0.1.1" description = "An integration package connecting Nomic and LangChain" authors = [] readme = "README.md" @@ -13,7 +13,7 @@ license = "MIT" [tool.poetry.dependencies] python = ">=3.8.1,<4.0" langchain-core = ">=0.1.46,<0.3" -nomic = "^3.0.12" +nomic = "^3.0.29" [tool.poetry.group.test] optional = true