AdvancedTopics: generate_embedings.py

File generate_embedings.py, 1.8 KB (added by 231040, 13 days ago)
Line 
1from sentence_transformers import SentenceTransformer
2import psycopg2
3
4model = SentenceTransformer('all-MiniLM-L6-v2')
5
6conn = psycopg2.connect(
7 host="localhost",
8 port="5432",
9 database="your_db_name",
10 user="your_db_user",
11 password="your_password"
12)
13
14cur = conn.cursor()
15
16query="""
17SELECT
18 s.service_id,
19 s.service_name,
20 sc.category_name,
21 s.price,
22 s.duration_minutes
23FROM service s
24JOIN service_category sc
25ON s.service_category_id=sc.service_category_id
26"""
27
28cur.execute(query)
29
30services=cur.fetchall()
31
32for service in services:
33
34 service_id=service[0]
35 name=service[1]
36 category=service[2]
37 price=service[3]
38 duration=service[4]
39
40 keywords=[]
41
42 name_lower=name.lower()
43 category_lower=category.lower()
44
45 if "facial" in name_lower or "skin" in category_lower:
46 keywords.extend([
47 "skin care",
48 "dry skin",
49 "hydration",
50 "face treatment",
51 "healthy skin"
52 ])
53
54 if "hair" in category_lower:
55 keywords.extend([
56 "hair styling",
57 "hair care",
58 "hair treatment"
59 ])
60
61 if "nail" in category_lower:
62 keywords.extend([
63 "manicure",
64 "pedicure",
65 "nail care"
66 ])
67
68 text=f"""
69 Service: {name}
70 Category: {category}
71 Price: {price}
72 Duration: {duration}
73 Keywords: {' '.join(keywords)}
74 """
75
76 embedding=model.encode(text).tolist()
77
78 cur.execute(
79 """
80 UPDATE service
81 SET embedding=%s
82 WHERE service_id=%s
83 """,
84 (embedding,service_id)
85 )
86
87conn.commit()
88
89cur.close()
90conn.close()
91
92print("Embeddings regenerated")