File size: 689 Bytes
ff4ec71 e25c291 ff4ec71 ca970dc ff4ec71 ca970dc 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 |
import pandas as pd
import streamlit as st
import torch
from sentence_transformers import SentenceTransformer
@st.cache(allow_output_mutation=True)
def load_model(model_name):
# Lazy downloading
model = SentenceTransformer(model_name)
return model
def load_embeddings():
# embedding pre-generated
corpus_emb = torch.load(
"./embeddings/descriptions_emb_200000_examples_github.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_200000_python_examples_github.csv")
return corpus_texts
|