File size: 696 Bytes
7e27e2f |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import duckdb
import polars as pl
from datasets import load_dataset
from model2vec import StaticModel
# Load a model from the HuggingFace hub (in this case the potion-base-8M model)
model_name = "minishlab/M2V_multilingual_output"
model = StaticModel.from_pretrained(model_name)
# Make embeddings
ds = load_dataset("fka/awesome-chatgpt-prompts")
df = ds["train"].to_polars()
embeddings = model.encode(df["prompt"])
df = df.with_columns(pl.Series(embeddings).alias("embeddings"))
vector = model.encode("vector search", show_progress_bar=True)
duckdb.sql(
query=f"""
SELECT *
FROM df
ORDER BY array_distance(embeddings, {vector.tolist()}::FLOAT[256])
LIMIT 1
"""
).show()
|