import pandas as pd
from sentence_transformers import SentenceTransformer
import streamlit as st
import torch


@st.cache(allow_output_mutation=True)
def load_model(model_name):
    # Lazy downloading
    model = SentenceTransformer(model_name)
    return model


@st.cache(allow_output_mutation=True)
def load_embeddings():
    # embedding pre-generated
    corpus_emb = torch.load(
        "./embeddings/descriptions_emb_100000_examples.pt",
        map_location=torch.device("cpu"),
    )
    return corpus_emb


@st.cache(allow_output_mutation=True)
def load_texts():
    # texts database pre-generated
    corpus_texts = pd.read_csv("./data/codesearchnet_100000_examples.csv")
    return corpus_texts