# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/45_nlp_interp.ipynb (unless otherwise specified).

__all__ = ['InterpEmbeddingsTokenizer']

# Cell
from ..interp.latent import InterpEmbeddings
import pandas as pd

# Cell
class InterpEmbeddingsTokenizer(InterpEmbeddings):
    def __init__(self,
                 embedding_matrix,
                 tokenizer):
        """
        embedding_matrix: np.ndarray, embedding matrix of shape:
            (num_items, hidden_size)
        tokenizer: a huggingface tokenizer
        """
        super().__init__(
            embedding_matrix,
            dict((v, k) for k, v in tokenizer.vocab.items()))
        self.tokenizer = tokenizer

    def search(
        self,
        word: str,
        filter_special_token: bool = True,
        top_k: int = 20,
    ) -> pd.DataFrame:
        """
        search for similar words with embedding and
            tokenizer's encode/ decode
        """
        token_ids = self.tokenizer.encode(word)
        if filter_special_token:
            token_ids = list(t for t in token_ids if t > 110)

        # combine multiple tokens into 1
        vec = self.base[token_ids].mean(0)

        # distance search
        closest, similarity = self.cosine.search(vec, return_similarity=True)
        tokens = self.tokenizer.convert_ids_to_tokens(closest)
        return pd.DataFrame({
            "tokens": tokens,
            "idx": closest,
            "similarity": similarity}).head(top_k)