Advanced: generate_embeddings.py

File generate_embeddings.py, 2.5 KB (added by 231166, 2 days ago)
Line 
1import psycopg2
2from sentence_transformers import SentenceTransformer
3
4# 1. povrzuvanje so vidi_db
5try:
6 conn = psycopg2.connect(
7 dbname="vidi_db",
8 user="postgres",
9 password="...",
10 host="localhost",
11 port="5432"
12 )
13 cursor = conn.cursor()
14 print("Успешно поврзување со локалната vidi_db!")
15except Exception as e:
16 print(f"Грешка при поврзување со базата: {e}")
17 exit()
18
19# 2. Vcituvanje na NLP modelot za gen vektori
20print("Се вчитува моделот за embeddings...")
21model = SentenceTransformer('all-MiniLM-L6-v2')
22
23try:
24 print("Ги влечам и ги групирам коментарите по содржина од локалната база...")
25 query = """
26 SELECT r.ContentContentID, m.title, STRING_AGG(r.Comment, ' ') as all_reviews
27 FROM "review" r
28 JOIN "media" m ON m.contentid = r.ContentContentID
29 GROUP BY r.ContentContentID, m.title;
30 """
31 cursor.execute(query)
32 rows = cursor.fetchall()
33
34 if not rows:
35 print("Нема нови филмови со коментари кои чекаат за вектор.")
36 else:
37 print(f"Пронајдени се {len(rows)} филмови за обработка на коментари...")
38
39 updated_count = 0
40 for row in rows:
41 content_id = row[0]
42 title = row[1]
43 all_reviews = row[2] if row[2] else ""
44
45 combined_context = f"{title}. Reviews: {all_reviews}"
46
47 embedding_vector = model.encode(combined_context).tolist()
48 truncated_vector = embedding_vector[:4] # Сечеме на 4 димензии за vector(4)
49
50 cursor.execute(
51 'UPDATE "media" SET embedding = %s WHERE contentid = %s;',
52 (truncated_vector, content_id)
53 )
54 updated_count += 1
55
56 if updated_count % 100 == 0:
57 print(f" Процесирани {updated_count} филмови локално...")
58
59 conn.commit()
60 print(f"\nУспешно ажурирани {updated_count} записи во локалната табела Media!")
61
62except Exception as e:
63 print(f"Настана грешка при работа со локалната база: {e}")
64 conn.rollback()
65finally:
66 cursor.close()
67 conn.close()
68 print("Врската со локалниот Postgres е затворена.")