import duckdb import lancedb from src.setting import AVAILABLE_WORDS class VectorDatabaseHandler: QUERY_TEMPLATE = "SELECT word, vector FROM {table_name} WHERE word = '{user_word}'" def __init__(self, db_path: str, table_name: str, metrics_cfg: dict): db = lancedb.connect(db_path) self.metrics_cfg = metrics_cfg self.embeddings_tbl = db.open_table(table_name) def __call__(self, guessed_word: str, supposed_word: str) -> dict: arrow_table = self.embeddings_tbl.to_arrow() word_embedding = self.get_word_vector(guessed_word, "arrow_table") df_emb = self.embeddings_tbl.search(word_embedding) \ .metric(self.metrics_cfg.metric) \ .limit(len(AVAILABLE_WORDS)) \ .to_df() supposed_word_row = df_emb[df_emb['word'] == supposed_word].iloc[0] cosine_distance = supposed_word_row['_distance'] words_between_count = len(df_emb[df_emb['_distance'] < cosine_distance]) closest_word = df_emb[df_emb['word'] != guessed_word].iloc[0]['word'] if words_between_count else supposed_word return { "score": cosine_distance, "rating": words_between_count, "percentage": 100 - words_between_count / len(df_emb) * 100, "closest_word": closest_word } def get_word_vector(self, word: str, table_name: str): vector = duckdb.query( self.QUERY_TEMPLATE.format(table_name=table_name, user_word=word) ).to_df()["vector"].values[0] return vector