AdvancedTopics: embed_task_requests_description.py

from sentence_transformers import SentenceTransformer
import psycopg2

model = SentenceTransformer("all-MiniLM-L6-v2")

conn = psycopg2.connect(
    host="localhost",
    database="tcs_local",
    user="postgres",
    password="postgres123"
)

read_cursor = conn.cursor()
write_cursor = conn.cursor()

try:
    read_cursor.execute("""
        SELECT
            tr.id,
            tr.description
        FROM
            TaskRequest tr

        LEFT JOIN
            task_request_description_vector trdv
            ON trdv.task_request_id = tr.id

        WHERE
            tr.status = 'OPEN'
            AND tr.description IS NOT NULL
            AND trdv.task_request_id IS NULL

        ORDER BY
            tr.id
    """)


    for task_request_id, description in read_cursor.fetchall():

        embedding = model.encode(
            description,
            convert_to_numpy=True,
            normalize_embeddings=True
        ).tolist()

        write_cursor.execute("""
            INSERT INTO task_request_description_vector
            (
                task_request_id,
                embedding
            )
            VALUES
            (%s, %s)

            ON CONFLICT (task_request_id)
            DO UPDATE SET
                embedding = EXCLUDED.embedding
        """, (task_request_id, embedding))

    conn.commit()

except Exception as e:
    conn.rollback()
    print(e)

finally:
    read_cursor.close()
    write_cursor.close()
    conn.close()