Added OpenSSL in linux arm build step in ci.yaml. Bumped to 0.1.9

Author: wdoppenberg

.github/workflows/ci.yml

@@ -163,6 +163,9 @@ jobs:
         with:
           python-version: ${{ matrix.py-version }}
 
+      - name: Install OpenSSL
+        run: sudo apt-get install libssl-dev
+
       - name: Build wheels
         uses: PyO3/maturin-action@v1
         with:

Cargo.lock

@@ -1,6 +1,6 @@
 [package]
 name = "polars-candle"
-version = "0.1.7"
+version = "0.1.9"
 edition = "2021"
 authors = ["Wouter Doppenberg <[email protected]>"]
 description = "A text embedding extension for the Polars Dataframe library."

Cargo.toml

@@ -6,6 +6,42 @@
 
 ModelName = Literal["bert-base-uncased", "bert-base-cased"]
 
+# Old models that don't belong to any organization
+# Source https://github.com/UKPLab/sentence-transformers/blob/c0fc0e8238f7f48a1e92dc90f6f96c86f69f1e02/sentence_transformers/SentenceTransformer.py#L219
+__MODEL_HUB_ORGANIZATION__ = "sentence-transformers"
+__BASIC_TRANSFORMER_MODELS__ = (
+    "bert-base-cased-finetuned-mrpc",
+    "bert-base-cased",
+    "bert-base-chinese",
+    "bert-base-german-cased",
+    "bert-base-german-dbmdz-cased",
+    "bert-base-german-dbmdz-uncased",
+    "bert-base-multilingual-cased",
+    "bert-base-multilingual-uncased",
+    "bert-base-uncased",
+    "bert-large-cased-whole-word-masking-finetuned-squad",
+    "bert-large-cased-whole-word-masking",
+    "bert-large-cased",
+    "bert-large-uncased-whole-word-masking-finetuned-squad",
+    "bert-large-uncased-whole-word-masking",
+    "bert-large-uncased",
+    "camembert-base",
+    "distilbert-base-cased-distilled-squad",
+    "distilbert-base-cased",
+    "distilbert-base-german-cased",
+    "distilbert-base-multilingual-cased",
+    "distilbert-base-uncased-distilled-squad",
+    "distilbert-base-uncased-finetuned-sst-2-english",
+    "distilbert-base-uncased",
+    "distilgpt2",
+    "distilroberta-base",
+    "roberta-base-openai-detector",
+    "roberta-base",
+    "roberta-large-mnli",
+    "roberta-large-openai-detector",
+    "roberta-large",
+)
+
 
 @pl.api.register_expr_namespace("candle")
 class CandleExt:
@@ -23,7 +59,7 @@ class CandleExt:
     ...     ],
     ... })
     >>> df.with_columns(
-    ...     pl.col("text").candle.embed_text("bert-base-uncased")
+    ...     pl.col("text").candle.embed_text("bert-base-cased")
     ... )
     """
 
@@ -57,6 +93,8 @@ def embed_text(
         Expr
                 An expression with the embedded text.
         """
+        if "/" not in model_repo:
+            model_repo = __MODEL_HUB_ORGANIZATION__ + "/" + model_repo
 
         return register_plugin_function(
             plugin_path=Path(__file__).parent,

polars_candle/candle_ext.py

@@ -1,6 +1,6 @@
 [project]
 name = "polars_candle"
-version = "0.1.7"
+version = "0.1.9"
 requires-python = ">=3.9"
 license = "Apache-2.0"
 classifiers = [
@@ -14,7 +14,7 @@ keywords = ["polars", "dataframe", "embedding", "nlp", "candle"]
 
 [tool.poetry]
 name = "polars_candle"
-version = "0.1.7"
+version = "0.1.8"
 description = "A text embedding extension for the Polars Dataframe library."
 authors = ["Wouter Doppenberg <[email protected]>"]
 

pyproject.toml

@@ -42,6 +42,16 @@ def test_basic_two_sentences_with_gpu():
     )
 
 
+def test_basic_model_no_repo():
+    df = pl.DataFrame({"s": ["This is a sentence", "This is another sentence"]})
+
+    df = df.with_columns(
+        pl.col("s")
+        .candle.embed_text("nli-distilbert-base", device="gpu")
+        .alias("s_embedding")
+    )
+
+
 def test_basic_with_none():
     df = pl.DataFrame(
         {"s": ["This is a sentence", None, "This is another sentence", None]}