Skip to content

Make sentence_transformers an optional dependency #674

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions examples/amazon_s3_embedding/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""
This example uses SentenceTransformerEmbed which requires the optional sentence-transformers dependency.

Install with:
pip install 'cocoindex[sentence-transformers]'

Or install the dependency directly:
pip install sentence-transformers
"""

from dotenv import load_dotenv
from psycopg_pool import ConnectionPool
import cocoindex
Expand Down
10 changes: 10 additions & 0 deletions examples/code_embedding/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""
This example uses SentenceTransformerEmbed which requires the optional sentence-transformers dependency.

Install with:
pip install 'cocoindex[sentence-transformers]'

Or install the dependency directly:
pip install sentence-transformers
"""

from dotenv import load_dotenv
from psycopg_pool import ConnectionPool
from pgvector.psycopg import register_vector
Expand Down
10 changes: 10 additions & 0 deletions examples/fastapi_server_docker/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""
This example uses SentenceTransformerEmbed which requires the optional sentence-transformers dependency.

Install with:
pip install 'cocoindex[sentence-transformers]'

Or install the dependency directly:
pip install sentence-transformers
"""

import cocoindex
import uvicorn
from dotenv import load_dotenv
Expand Down
10 changes: 10 additions & 0 deletions examples/gdrive_text_embedding/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""
This example uses SentenceTransformerEmbed which requires the optional sentence-transformers dependency.

Install with:
pip install 'cocoindex[sentence-transformers]'

Or install the dependency directly:
pip install sentence-transformers
"""

from dotenv import load_dotenv
from psycopg_pool import ConnectionPool
import cocoindex
Expand Down
10 changes: 10 additions & 0 deletions examples/pdf_embedding/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""
This example uses SentenceTransformerEmbed which requires the optional sentence-transformers dependency.

Install with:
pip install 'cocoindex[sentence-transformers]'

Or install the dependency directly:
pip install sentence-transformers
"""

import cocoindex
import os
import tempfile
Expand Down
10 changes: 10 additions & 0 deletions examples/text_embedding/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""
This example uses SentenceTransformerEmbed which requires the optional sentence-transformers dependency.

Install with:
pip install 'cocoindex[sentence-transformers]'

Or install the dependency directly:
pip install sentence-transformers
"""

from dotenv import load_dotenv
from psycopg_pool import ConnectionPool
from pgvector.psycopg import register_vector
Expand Down
10 changes: 10 additions & 0 deletions examples/text_embedding_qdrant/main.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,13 @@
"""
This example uses SentenceTransformerEmbed which requires the optional sentence-transformers dependency.

Install with:
pip install 'cocoindex[sentence-transformers]'

Or install the dependency directly:
pip install sentence-transformers
"""

from dotenv import load_dotenv
from qdrant_client import QdrantClient
import cocoindex
Expand Down
2 changes: 1 addition & 1 deletion pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,6 @@ authors = [{ name = "CocoIndex", email = "cocoindex.io@gmail.com" }]
readme = "README.md"
requires-python = ">=3.11"
dependencies = [
"sentence-transformers>=3.3.1",
"click>=8.1.8",
"rich>=14.0.0",
"python-dotenv>=1.1.0",
Expand All @@ -31,6 +30,7 @@ features = ["pyo3/extension-module"]
[project.optional-dependencies]
test = ["pytest"]
dev = ["ruff", "pre-commit"]
sentence-transformers = ["sentence-transformers>=3.3.1"]

[tool.mypy]
python_version = "3.11"
Expand Down
20 changes: 20 additions & 0 deletions python/cocoindex/functions.py
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,14 @@
if TYPE_CHECKING:
import sentence_transformers

# Check if sentence_transformers is available
try:
import sentence_transformers

_SENTENCE_TRANSFORMERS_AVAILABLE = True
except ImportError:
_SENTENCE_TRANSFORMERS_AVAILABLE = False


class ParseJson(op.FunctionSpec):
"""Parse a text into a JSON object."""
Expand Down Expand Up @@ -58,6 +66,10 @@ class SentenceTransformerEmbed(op.FunctionSpec):

model: The name of the SentenceTransformer model to use.
args: Additional arguments to pass to the SentenceTransformer constructor. e.g. {"trust_remote_code": True}

Note:
This function requires the optional sentence-transformers dependency.
Install it with: pip install 'cocoindex[sentence-transformers]'
"""

model: str
Expand All @@ -72,6 +84,14 @@ class SentenceTransformerEmbedExecutor:
_model: "sentence_transformers.SentenceTransformer"

def analyze(self, text: Any) -> type:
if not _SENTENCE_TRANSFORMERS_AVAILABLE:
raise ImportError(
"sentence_transformers is required for SentenceTransformerEmbed function. "
"Install it with one of these commands:\n"
" pip install 'cocoindex[sentence-transformers]'\n"
" pip install sentence-transformers"
)

import sentence_transformers # pylint: disable=import-outside-toplevel

args = self.spec.args or {}
Expand Down