https://github.com/lmmx/polars-sbert
A Polars plugin for embedding data with Sentence Transformers
https://github.com/lmmx/polars-sbert
Last synced: about 2 months ago
JSON representation
A Polars plugin for embedding data with Sentence Transformers
- Host: GitHub
- URL: https://github.com/lmmx/polars-sbert
- Owner: lmmx
- Created: 2025-02-21T00:00:33.000Z (over 1 year ago)
- Default Branch: master
- Last Pushed: 2025-02-21T00:00:41.000Z (over 1 year ago)
- Last Synced: 2025-10-05T06:22:51.907Z (about 1 year ago)
- Language: Python
- Size: 72.3 KB
- Stars: 1
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
- awesome-polars - polars-sbert - Polars plugin for embedding data with Sentence Transformers by [@lmmx](https://github.com/lmmx). (Libraries/Packages/Scripts / Polars plugins)
README
# Polars SBERT
A Polars plugin for using embeddings with Sentence Transformers (SBERT).
## Installation
Install the full Polars extra and CUDA 12.4-based PyTorch (`cu124`):
```bash
pip install "polars-sentence-transformers[polars,cu124]"
```
On older CPUs install `polars-lts-cpu`
```bash
pip install polars-sentence-transformers[polars-lts-cpu]
```
and for CPU-only PyTorch use the `cpu-only` extra:
```bash
pip install polars-sentence-transformers[cpu-only]
```
## Features
- Embed from a DataFrame by specifying the source column(s)
- Re-order/filter rows by semantic similarity to a query
- Efficiently reuse loaded models via a global registry (no repeated model loads)
## Usage
TBD
## Demo
See [demo.py](https://github.com/lmmx/polars-sbert/tree/master/demo.py)
```
Corpus:
shape: (3, 2)
┌───────────────────────────┬───────────────────────────────┐
│ title ┆ body │
│ --- ┆ --- │
│ str ┆ str │
╞═══════════════════════════╪═══════════════════════════════╡
│ Living With Normal Pets ┆ Pet owner's manual │
│ My Cat Wants To Murder Me ┆ Scary stories about evil cats │
│ A Guide to Walking Dogs ┆ Learning to lead │
└───────────────────────────┴───────────────────────────────┘
Embedded title and body columns with TaylorAI/gte-tiny
Retrieve top k=None rows with query 'horror stories'
shape: (3, 4)
┌───────────────────────────┬────────────────────────────┬────────────────────────────┬────────────┐
│ title ┆ body ┆ embedding ┆ similarity │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ list[f32] ┆ f64 │
╞═══════════════════════════╪════════════════════════════╪════════════════════════════╪════════════╡
│ My Cat Wants To Murder Me ┆ Scary stories about evil ┆ [-0.1195, -0.289812, … ┆ 0.858065 │
│ ┆ cats ┆ 0.14939… ┆ │
│ Living With Normal Pets ┆ Pet owner's manual ┆ [-0.168101, -0.158491, … ┆ 0.754353 │
│ ┆ ┆ 0.488… ┆ │
│ A Guide to Walking Dogs ┆ Learning to lead ┆ [-0.388657, -0.37131, … ┆ 0.727875 │
│ ┆ ┆ 0.6422… ┆ │
└───────────────────────────┴────────────────────────────┴────────────────────────────┴────────────┘
Currently loaded models: ['TaylorAI/gte-tiny']
Cleared model registry
Currently loaded models: []
```
## Contributing
Feel free to open issues or submit pull requests for improvements or bug fixes.
## License
MIT License