File size: 706 Bytes
ff4ec71 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
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
|