import os

import pandas as pd
import lancedb
from lancedb.embeddings import with_embeddings
from sentence_transformers import SentenceTransformer

from setting import CFG, AVAILABLE_WORDS


df = pd.DataFrame(AVAILABLE_WORDS, columns=['word'])

model = SentenceTransformer(CFG.model.name)

data = with_embeddings(
    func=lambda texts: model.encode(texts),
    data=df, column="word", show_progress=True
)

if not os.path.exists(CFG.db.lance_db_folder_path):
    os.makedirs(CFG.db.lance_db_folder_path)

db = lancedb.connect(CFG.db.lance_db_folder_path)
table = db.create_table(CFG.db.table_name, data)
print("Table created")