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>

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"]
 

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"

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