Test-API / src /create-embedding.py
lpetrl's picture
feat(API): Implemented basic functionality.
94e8fb8 verified
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")