--

Author: wrmsr

ommlx/__about__.py

@@ -13,17 +13,23 @@ class Project(ProjectBase):
     ]
 
     optional_dependencies = {
+        'backends': [
+            'llama-cpp-python ~= 0.2',
+
+            'openai ~= 1.47',
+
+            'sentence_transformers ~= 3.1',
+
+            'transformers ~= 4.44',
+        ],
+
         'huggingface': [
             'huggingface-hub ~= 0.25',
             'datasets ~= 3.0',
         ],
 
-        'llamacpp': [
-            'llama-cpp-python ~= 0.2',
-        ],
-
-        'openai': [
-            'openai ~= 1.47',
+        'pillow': [
+            'pillow ~= 10.4',
         ],
 
         'tinygrad': [
@@ -35,10 +41,6 @@ class Project(ProjectBase):
             'torch ~= 2.4',
         ],
 
-        'transformers': [
-            'transformers ~= 4.44',
-        ],
-
         'wiki': [
             'mwparserfromhell ~= 0.6',
             'wikitextparser ~= 0.56',  # !! GPL

ommlx/minichain/backends/openai.py

@@ -10,6 +10,7 @@
 from ..chat import SystemMessage
 from ..chat import ToolExecutionResultMessage
 from ..chat import UserMessage
+from ..content import Content
 from ..content import Text
 from ..embeddings import Embedding
 from ..embeddings import EmbeddingModel_
@@ -87,10 +88,10 @@ def generate(self, request: Request[ChatRequest]) -> Response[AiMessage]:
 class OpenaiEmbeddingModel(EmbeddingModel_):
     model = 'text-embedding-3-small'
 
-    def generate(self, request: Request[str]) -> Response[Embedding]:
+    def generate(self, request: Request[Content]) -> Response[Embedding]:
         response = openai.embeddings.create(
             model=self.model,
-            input=request.v,
+            input=check.isinstance(request.v, Text).s,
         )
 
         return Response(Embedding(response.data[0].embedding))

ommlx/minichain/backends/sentencetransformers.py

@@ -0,0 +1,30 @@
+import typing as ta
+
+from omlish import check
+from omlish import lang
+
+from ..content import Content
+from ..content import Text
+from ..embeddings import Embedding
+from ..embeddings import EmbeddingModel_
+from ..models import Request
+from ..models import Response
+
+
+if ta.TYPE_CHECKING:
+    import PIL.Image as pimg  # noqa
+    import sentence_transformers as st
+else:
+    pimg = lang.proxy_import('PIL.Image')
+    st = lang.proxy_import('sentence_transformers')
+
+
+class SentencetransformersEmbeddingModel(EmbeddingModel_):
+    model = 'clip-ViT-B-32'
+
+    def generate(self, request: Request[Content]) -> Response[Embedding]:
+        mdl = st.SentenceTransformer(self.model)
+
+        response = mdl.encode(check.isinstance(request.v, Text).s)
+
+        return Response(Embedding(response.tolist()))

ommlx/minichain/cli/main.py

@@ -14,6 +14,7 @@
 from ..backends.openai import OpenaiChatModel
 from ..backends.openai import OpenaiEmbeddingModel
 from ..backends.openai import OpenaiPromptModel
+from ..backends.sentencetransformers import SentencetransformersEmbeddingModel
 from ..backends.transformers import TransformersPromptModel
 from ..chat import Chat
 from ..chat import ChatModel
@@ -139,6 +140,7 @@ def _run_prompt(
 
 EMBEDDING_MODEL_BACKENDS: ta.Mapping[str, type[EmbeddingModel]] = {
     'openai': OpenaiEmbeddingModel,
+    'sentencetransformers': SentencetransformersEmbeddingModel,
 }
 
 
@@ -149,7 +151,7 @@ def _run_embed(
 ) -> None:
     mdl = EMBEDDING_MODEL_BACKENDS[backend or DEFAULT_BACKEND]()
 
-    response = mdl.generate(Request(prompt))
+    response = mdl.generate(Request(Text(prompt)))
 
     print(json.dumps_compact(response.v.v))
 

ommlx/minichain/embeddings.py

@@ -4,6 +4,7 @@
 
 from omlish import lang
 
+from .content import Content
 from .models import Model
 from .models import Request
 from .models import Response
@@ -19,10 +20,10 @@ class Embedding:
     v: ta.Sequence[float]
 
 
-EmbeddingModel: ta.TypeAlias = Model[str, Embedding]
+EmbeddingModel: ta.TypeAlias = Model[Content, Embedding]
 
 
 class EmbeddingModel_(EmbeddingModel, lang.Abstract):  # noqa
     @abc.abstractmethod
-    def generate(self, request: Request[str]) -> Response[Embedding]:
+    def generate(self, request: Request[Content]) -> Response[Embedding]:
         raise NotImplementedError

ommlx/tests/__init__.py

@@ -158,7 +158,7 @@ einops==0.8.0
 
 ### infra
 
-boto3>=1.35.24
+boto3>=1.35.2G4
 
 google-cloud-compute==1.19.2
 
@@ -269,3 +269,5 @@ usearch==2.15.1
 llama-index-core==0.11.10
 
 unstructured==0.15.13
+
+sentence_transformers==3.1.1