/
rnekrasov
/
pgvector-python
Обзор
Документация
Войти
/
rnekrasov
/
pgvector-python
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
examples/sentence_embeddings.py
28 строк
995 B
Andrew Kane
Updated example [skip ci]
09 сен 2023, 02:15
09 сен 2023, 02:15
8bc6a3d
Код
Авторство
О чём код?
from pgvector.psycopg import register_vector import psycopg from sentence_transformers import SentenceTransformer conn = psycopg.connect(dbname='pgvector_example', autocommit=True) conn.execute('CREATE EXTENSION IF NOT EXISTS vector') register_vector(conn) conn.execute('DROP TABLE IF EXISTS documents') conn.execute('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))') input = [ 'The dog is barking', 'The cat is purring', 'The bear is growling' ] model = SentenceTransformer('all-MiniLM-L6-v2') embeddings = model.encode(input) for content, embedding in zip(input, embeddings): conn.execute('INSERT INTO documents (content, embedding) VALUES (%s, %s)', (content, embedding)) document_id = 1 neighbors = conn.execute('SELECT content FROM documents WHERE id != %(id)s ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = %(id)s) LIMIT 5', {'id': document_id}).fetchall() for neighbor in neighbors: print(neighbor[0])