mirror of
https://github.com/hwchase17/langchain
synced 2024-11-10 01:10:59 +00:00
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
|
import os
|
||
|
|
from typing import Any, Dict, List
|
||
|
|
|
||
|
|
import together # type: ignore
|
||
|
|
from langchain_core.embeddings import Embeddings
|
||
|
|
from langchain_core.pydantic_v1 import BaseModel, SecretStr, root_validator
|
||
|
|
from langchain_core.utils import convert_to_secret_str
|
||
|
|
|
||
|
|
|
||
|
|
class TogetherEmbeddings(BaseModel, Embeddings):
|
||
|
|
"""TogetherEmbeddings embedding model.
|
||
|
|
|
||
|
|
Example:
|
||
|
|
.. code-block:: python
|
||
|
|
|
||
|
|
from langchain_together import TogetherEmbeddings
|
||
|
|
|
||
|
|
model = TogetherEmbeddings(
|
||
|
|
model='togethercomputer/m2-bert-80M-8k-retrieval'
|
||
|
|
)
|
||
|
|
"""
|
||
|
|
|
||
|
|
_client: together.Together
|
||
|
|
together_api_key: SecretStr = convert_to_secret_str("")
|
||
|
|
model: str
|
||
|
|
|
||
|
|
@root_validator()
|
||
|
|
def validate_environment(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
||
|
|
"""Validate environment variables."""
|
||
|
|
together_api_key = convert_to_secret_str(
|
||
|
|
values.get("together_api_key") or os.getenv("TOGETHER_API_KEY") or ""
|
||
|
|
)
|
||
|
|
values["together_api_key"] = together_api_key
|
||
|
|
|
||
|
|
# note this sets it globally for module
|
||
|
|
# there isn't currently a way to pass it into client
|
||
|
|
together.api_key = together_api_key.get_secret_value()
|
||
|
|
values["_client"] = together.Together()
|
||
|
|
return values
|
||
|
|
|
||
|
|
def embed_documents(self, texts: List[str]) -> List[List[float]]:
|
||
|
|
"""Embed search docs."""
|
||
|
|
return [
|
||
|
|
i.embedding
|
||
|
|
for i in self._client.embeddings.create(input=texts, model=self.model).data
|
||
|
|
]
|
||
|
|
|
||
|
|
def embed_query(self, text: str) -> List[float]:
|
||
|
|
"""Embed query text."""
|
||
|
|
return self.embed_documents([text])[0]
|