vectorizer.py

import gensim
import gensim.downloader
from gensim.models import KeyedVectors
import numpy as np
import pandas as pd
import os
from supabase import acreate_client, AsyncClient
from dotenv import load_dotenv

class Vectorizer:
    """
    A class to:
        - Generate embeddings of words
        - Query for words from Supabase database based on vector similarity
        - Return matching ASL videos for words
    """

    def load_kv(self, model_name='word2vec-google-news-300'):
        """
        Returns a KeyedVector object loaded from gensim
        """
        model_path = os.path.join(os.getcwd(), 'gensim-data', 'GoogleNews-vectors-negative300.bin.gz')
        try:
            print(f"Loading model from {model_path}")
            kv = KeyedVectors.load_word2vec_format(model_path, binary=True)
            print("Word2Vec model loaded successfully as KeyedVectors object.")
            return kv
        except FileNotFoundError:
            print(f"Error: Model file not found at {model_path}. Trying to download...")
            kv = gensim.downloader.load(model_name)  # returns a keyedvector
            print("Word2Vec model loaded successfully as KeyedVectors object.")
            return kv
        except Exception as e:
            print(f"Unable to load embedding model from gensim: {e}")
            return None

    async def initialize_supabase(self):
        url: str = os.environ.get("SUPABASE_URL")
        key: str = os.environ.get("SUPABASE_KEY")
        supabase: AsyncClient = await acreate_client(url, key)
        return supabase

    def __init__(self):
        load_dotenv()
        self.kv = self.load_kv()
        self.supabase = None  # Will be initialized when needed

    async def ensure_supabase_initialized(self):
        """Ensure Supabase client is initialized"""
        if self.supabase is None:
            self.supabase = await self.initialize_supabase()

    def encode(self, word):
        print(f"encoding {word}")
        if self.kv is None:
            print("KeyedVectors not loaded")
            return None
        if word in self.kv.key_to_index:
            return self.kv[word]
        else:
            print(f"Error: {word} is not in the KeyedVector's vocabulary")
            # Try to find closest match
            try:
                closest_matches = self.kv.most_similar(word, topn=3)
                if closest_matches:
                    closest_word = closest_matches[0][0]
                    print(f"Using closest match '{closest_word}' for '{word}'")
                    return self.kv[closest_word]
                else:
                    print(f"No similar words found for '{word}'")
            except Exception as e:
                print(f"Error finding similar words: {e}")
            return None
    
    def encode_and_format(self, word):
        """
        Apply encoding function to each word.
        Prettify the encoding to match expected format for Supabase vectors
        """
        enc = self.encode(word)
        return "[" + ",".join(map(str, enc.tolist())) + "]" if enc is not None else None

    async def vector_query_from_supabase(self, query):
        try:
            await self.ensure_supabase_initialized()
            query_embedding = self.encode(query)
            
            if query_embedding is None:
                return {
                    "match": False,
                    "error": f"'{query}' not in vocabulary and no similar words found"
                }

            query_embedding = query_embedding.tolist()

            if self.supabase is not None:
                result = await self.supabase.rpc(
                    "match_vector",
                    {
                        "query_embedding": query_embedding,
                        "match_threshold": 0.0,
                        "match_count": 1
                    }
                ).execute()

                data = result.data
                if data:
                    match = data[0]
                    return {
                        "match": True,
                        "query": query,
                        "matching_word": match["word"],
                        "video_url": match["video_url"],
                        "similarity": match["similarity"]
                    }
                else:
                    return {"match": False}
            else:
                return {"match": False, "error": "Supabase not initialized"}

        except Exception as e:
            print(f"RPC call failed: {e}")
            return {"match": False, "error": str(e)}


def load_filtered_kv(model_name='word2vec-google-news-300', vocab=None):
    """
    Returns a KeyedVector object whose vocabulary
    consists of the words in vocab
    """
    if vocab is None:
        vocab = []
    try:
        # gensim.downloader.load returns a KeyedVector
        original_kv = gensim.downloader.load(model_name)
        if vocab:
            filtered_key2vec_map = {}
            for key in vocab:
                if key in original_kv.key_to_index:
                    filtered_key2vec_map[key] = original_kv[key]

            new_kv = gensim.models.KeyedVectors(
                vector_size=original_kv.vector_size)
            new_kv.add_vectors(list(filtered_key2vec_map.keys()),
                               np.array(list(filtered_key2vec_map.values())))
            return original_kv
        else:
            return original_kv
    except Exception as e:
        print(f"Unable to load embedding model from gensim: {e}")
        return None


async def main():
    vectorizer = Vectorizer()

    # Test exact word match
    vector = vectorizer.encode("test")
    print(vector)
    
    # Test words not in vocabulary with closest match fallback
    result = await vectorizer.vector_query_from_supabase("dog")
    print(result)
    result = await vectorizer.vector_query_from_supabase("cat")
    print(result)

    # read word list
    # df = pd.read_csv('videos_rows.csv')
    
    # # Add embeddings column - apply encode to each word
    # df['embedding'] = df['word'].apply(vectorizer.encode_and_format)

    # # Drop any rows that don't have an embedding
    # df = df.dropna(subset=['embedding'])
    # print(df.head())

    # df.to_csv("vectors.csv", index=False, columns=["word", "video_url", "embedding"], header=True)


if __name__ == "__main__":
    import asyncio
    asyncio.run(main())