content: import warnings
  
  import numpy as np
  from typing import Union, List
  from sklearn.metrics import pairwise_distances
  from gensim.models import KeyedVectors
  
  
  from whatlies.embedding import Embedding
  from whatlies.embeddingset import EmbeddingSet
  from whatlies.language._common import SklearnTransformerMixin
  
  
  class GensimLanguage(SklearnTransformerMixin):
      """
      This object is used to lazily fetch [Embedding][whatlies.embedding.Embedding]s or
      [EmbeddingSet][whatlies.embeddingset.EmbeddingSet]s from a keyed vector file.
      These files are generated by [gensim](https://radimrehurek.com/gensim/models/word2vec.html).
      This object is meant for retreival, not plotting.
  
      Important:
          The vectors are not given by this library they must be download/created upfront.
          A potential benefit of this is that you can train your own embeddings using
          gensim and visualise them using this library.
  
          Here's a snippet that you can use to train your own (very limited) word2vec embeddings.
  
          ```
          from gensim.test.utils import common_texts
          from gensim.models import Word2Vec
          model = Word2Vec(common_texts, size=10, window=5, min_count=1, workers=4)
          model.wv.save("wordvectors.kv")
          ```
  
          You can also download pre-trained embeddings that are hosted by the gensim project.
  
          ```
          import gensim.downloader as api
          # To check what models are available
          api.info()['models'].keys()
          # To download the vectors
          wv = api.load('glove-twitter-25')
          # This is typically saved in `~/gensim/data` but you can also edit these
          # vectors and save them someplace else if you'd like.
          wv.save("glove-twitter-25.kv")
          ```
  
          Note that if a word is not available in the keyed vectors file then we'll assume
          a zero vector. If you pass a sentence then we'll add together the embeddings vectors
          of the seperate words.
  
      Arguments:
          keyedfile: name of the model to load, be sure that it's downloaded or trained beforehand
  
      **Usage**:
  
      ```python
      > from whatlies.language import GensimLanguage
      > lang = GensimLanguage("wordvectors.kv")
      > lang['computer']
      > lang = GensimLanguage("wordvectors.kv")
      > lang[['computer', 'human', 'dog']]
      ```
      """
  
      def __init__(self, keyedfile):
          self.keyedfile = keyedfile
          self.kv = KeyedVectors.load(keyedfile)
  
      def __getitem__(self, query: Union[str, List[str]]):
          """
          Retreive a single embedding or a set of embeddings.
  
          Arguments:
              query: single string or list of strings
  
          **Usage**
          ```python
          > from whatlies.language import GensimLanguage
          > lang = GensimLanguage("wordvectors.kv")
          > lang['computer']
          > lang = GensimLanguage("wordvectors.kv")
          > lang[['computer', 'human', 'dog']]
          ```
          """
          if isinstance(query, str):
              if " " in query:
                  return Embedding(
                      query, np.sum([self[q].vector for q in query.split(" ")], axis=0)
                  )
              try:
                  vec = np.sum([self.kv[q] for q in query.split(" ")], axis=0)
              except KeyError:
                  vec = np.zeros(self.kv.vector_size)
              return Embedding(query, vec)
          return EmbeddingSet(*[self[tok] for tok in query])
  
      def _prepare_queries(self, lower):
          queries = [w for w in self.kv.vocab.keys()]
          if lower:
              queries = [w for w in queries if w.lower() == w]
          return queries
  
      def _calculate_distances(self, emb, queries, metric):
          vec = emb.vector
          vector_matrix = np.array([self[w].vector for w in queries])
          # there are NaNs returned, good to investigate later why that might be
          vector_matrix = np.array(
              [np.zeros(v.shape) if np.any(np.isnan(v)) else v for v in vector_matrix]
          )
          return pairwise_distances(vector_matrix, vec.reshape(1, -1), metric=metric)
  
      def score_similar(
          self,
          emb: Union[str, Embedding],
          n: int = 10,
          metric="cosine",
          lower=False,
      ) -> List:
          """
          Retreive a list of (Embedding, score) tuples that are the most similar to the passed query.
  
          Arguments:
              emb: query to use
              n: the number of items you'd like to see returned
              metric: metric to use to calculate distance, must be scipy or sklearn compatible
              lower: only fetch lower case tokens
  
          Returns:
              An list of ([Embedding][whatlies.embedding.Embedding], score) tuples.
          """
          if isinstance(emb, str):
              emb = self[emb]
  
          queries = self._prepare_queries(lower=lower)
          distances = self._calculate_distances(emb=emb, queries=queries, metric=metric)
          by_similarity = sorted(zip(queries, distances), key=lambda z: z[1])
  
          if len(queries) < n:
              warnings.warn(
                  f"We could only find {len(queries)} feasible words. Consider changing `top_n` or `lower`",
                  UserWarning,
              )
  
          return [(self[q], float(d)) for q, d in by_similarity[:n]]
  
      def embset_similar(
          self,
          emb: Union[str, Embedding],
          n: int = 10,
          lower=False,
          metric="cosine",
      ) -> EmbeddingSet:
          """
          Retreive an [EmbeddingSet][whatlies.embeddingset.EmbeddingSet] that are the most similar to the passed query.
  
          Arguments:
              emb: query to use
              n: the number of items you'd like to see returned
              metric: metric to use to calculate distance, must be scipy or sklearn compatible
              lower: only fetch lower case tokens
  
          Returns:
              An [EmbeddingSet][whatlies.embeddingset.EmbeddingSet] containing the similar embeddings.
          """
          embs = [
              w[0] for w in self.score_similar(emb=emb, n=n, lower=lower, metric=metric)
          ]
          return EmbeddingSet({w.name: w for w in embs})
  