/
rnekrasov
/
pgvector-python
Обзор
Документация
Войти
/
rnekrasov
/
pgvector-python
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
master
examples/openai_embeddings.py
29 строк
1 KB
Andrew Kane
Updated OpenAI model [skip ci]
26 янв 2024, 05:22
26 янв 2024, 05:22
2bb8c3c
Код
Авторство
О чём код?
from openai import OpenAI from pgvector.psycopg import register_vector import psycopg 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(1536))') input = [ 'The dog is barking', 'The cat is purring', 'The bear is growling' ] client = OpenAI() response = client.embeddings.create(input=input, model='text-embedding-3-small') embeddings = [v.embedding for v in response.data] 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])