tushifire's picture
Waits for the input first
365cc9e
import streamlit as st
import pickle
import pandas
from sentence_transformers import SentenceTransformer, util
import torch
st.title('Arxiv Paper Recommendation')
paper_you_like = st.text_input(
"Enter the title of any paper you like 👇",
placeholder = None
)
# @st.cache_resource
def get_sentences_data():
with open('sentences.pkl', 'rb') as f:
sentences = pickle.load(f)
return sentences
sentences = get_sentences_data()
# @st.cache_resource
def get_embeddings_data():
with open('embeddings.pkl', 'rb') as f:
embeddings = pickle.load(f)
return embeddings
embeddings = get_embeddings_data()
# @st.cache_resource
def get_model():
model = SentenceTransformer('all-MiniLM-L6-v2')
return model
model = get_model()
if paper_you_like: # if its not NONE
# Calculating the similarity between titles
cosine_scores = util.cos_sim(embeddings, model.encode(paper_you_like))
top_similar_papers = torch.topk(cosine_scores,dim=0, k=5,sorted=True)
# top_similar_papers
st.subheader('Recommended Papers are :scroll: ')
for i in top_similar_papers.indices:
st.write(sentences[i.item()])