espejelomar commited on
Commit
e25c291
1 Parent(s): 9ec02db

Update backend/utils.py

Browse files
Files changed (1) hide show
  1. backend/utils.py +3 -4
backend/utils.py CHANGED
@@ -1,7 +1,7 @@
1
  import pandas as pd
2
- from sentence_transformers import SentenceTransformer
3
  import streamlit as st
4
  import torch
 
5
 
6
 
7
  @st.cache(allow_output_mutation=True)
@@ -11,11 +11,10 @@ def load_model(model_name):
11
  return model
12
 
13
 
14
- @st.cache(allow_output_mutation=True)
15
  def load_embeddings():
16
  # embedding pre-generated
17
  corpus_emb = torch.load(
18
- "./embeddings/descriptions_emb_100000_examples.pt",
19
  map_location=torch.device("cpu"),
20
  )
21
  return corpus_emb
@@ -24,5 +23,5 @@ def load_embeddings():
24
  @st.cache(allow_output_mutation=True)
25
  def load_texts():
26
  # texts database pre-generated
27
- corpus_texts = pd.read_csv("./data/codesearchnet_100000_examples.csv")
28
  return corpus_texts
1
  import pandas as pd
 
2
  import streamlit as st
3
  import torch
4
+ from sentence_transformers import SentenceTransformer
5
 
6
 
7
  @st.cache(allow_output_mutation=True)
11
  return model
12
 
13
 
 
14
  def load_embeddings():
15
  # embedding pre-generated
16
  corpus_emb = torch.load(
17
+ "./embeddings/descriptions_emb_10000_examples_github.pt",
18
  map_location=torch.device("cpu"),
19
  )
20
  return corpus_emb
23
  @st.cache(allow_output_mutation=True)
24
  def load_texts():
25
  # texts database pre-generated
26
+ corpus_texts = pd.read_csv("./data/codesearchnet_10000_python_examples_github.csv")
27
  return corpus_texts