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