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:
Jared Van Bortel 2024-05-20 17:17:07 -04:00 committed by GitHub
parent 0e72ed39a0
commit 25d1c1c9bb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 56 additions and 10 deletions

View File

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

View File

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

View File

@ -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