HIPPOCRATES / embed.py
tbaro's picture
Upload 3 files
ab5694f verified
import pandas as pd
from sentence_transformers import SentenceTransformer
import sqlite3
import json
# Đường dẫn tệp CSV
csv_file_path = 'C:/Users/Thai Bao/Downloads/archive (1)/medquad.csv'
# Đọc dữ liệu
data = pd.read_csv(csv_file_path)
# Khởi tạo Sentence Transformer
model = SentenceTransformer('all-MiniLM-L6-v2')
# Xử lý giá trị thiếu (NaN) trong cột 'answer'
data['answer'] = data['answer'].fillna("")
# Embed cột "answer"
data['answer_vector'] = data['answer'].apply(lambda x: model.encode(x).tolist())
# Lưu dữ liệu và vector vào SQLite
db_path = 'answers.db'
conn = sqlite3.connect(db_path)
# Chuyển đổi embedding thành chuỗi JSON để lưu
data_to_save = data[['question', 'answer', 'source', 'focus_area', 'answer_vector']]
data_to_save['answer_vector'] = data_to_save['answer_vector'].apply(json.dumps)
# Lưu vào SQLite
data_to_save.to_sql('answers', conn, if_exists='replace', index=False)
conn.close()
# Khi đọc dữ liệu từ SQLite
conn = sqlite3.connect(db_path)
df = pd.read_sql('SELECT * FROM answers', conn)
df['answer_vector'] = df['answer_vector'].apply(json.loads) # Chuyển JSON về danh sách
conn.close()
# Hiển thị dữ liệu đọc lại
print(df.head())