# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['Format', 'SentenceGraph']

# %% ../nbs/00_core.ipynb 3
from sentence_transformers import SentenceTransformer
from sentence_transformers.util import cos_sim
from typing import List
from enum import Enum

# %% ../nbs/00_core.ipynb 4
class Format(Enum):
  Python = 0
  Torch = 1
  Numpy = 2
  Pandas = 3

# %% ../nbs/00_core.ipynb 5
class SentenceGraph:
    def __init__(self, model_name: str = None, model: None = None):
        if model_name == None or model_name == "" or model == None:
            self.model = model = SentenceTransformer('all-MiniLM-L6-v2')
        elif model is not None:
            self.model = model
        elif model_name is not None:
          self.model = model = SentenceTransformer(model_name)

    def createGraph(self, sentences: List[str], format: Format = Format.Python) -> List[List[float]]:
        # TODO: Wrap this all in a function to export a functional version of this as well.
        sentence_embeddings = self.model.encode(sentences)

        graph_len = len(sentence_embeddings)
        similarity_graph = [[0 for x in range(graph_len)] for y in range(graph_len)]

        for i, embeddingA in enumerate(sentence_embeddings):
          for j, embeddingB in enumerate(sentence_embeddings):
            similarity = cos_sim(embeddingA, embeddingB)

            # TODO when we get to python 3.10+ switch to pattern matching
            if format == Format.Python:
              similarity = similarity.tolist()[0][0]
            elif format == Format.Torch:
              similarity = similarity[0]
            elif format == Format.Pandas:
              # TODO: Implement eventual conversion to a dataframe
              similarity = similarity.tolist()[0][0]
            elif format == Format.Numpy:
              similarity = similarity[0].numpy()

            similarity_graph[i][j] = similarity
          
        return similarity_graph
