RecModel / app.py
1998Shubham's picture
Update app.py
0912f5a verified
raw
history blame
1.07 kB
import streamlit as st
import torch
import pandas as pd
from sentence_transformers import SentenceTransformer, util
import pickle
# Load data
data = pd.read_csv("./arxiv_data.csv")
titles = data["titles"]
# Load pre-trained SentenceTransformer model
model = SentenceTransformer("all-MiniLM-L6-v2")
# Load saved embeddings
with open("./embedding.pkl", "rb") as f:
Lencode = pickle.load(f)
# Load saved model
with open("./ModelRec.pkl", "rb") as f:
lModelRec = pickle.load(f)
def recomm(inputPaper):
encodePaper = lModelRec.encode(inputPaper)
cosine_score = util.cos_sim(Lencode, encodePaper)
top_scores = torch.topk(cosine_score, dim=0, k=4)
paperList = []
for i in top_scores.indices:
paperList.append(titles[i.item()])
return paperList
# Streamlit UI
st.title("Paper Recommendation System")
input_paper = st.text_input("Enter the name of the paper")
if st.button("Recommend"):
recommended_papers = recomm(input_paper)
st.write("Recommended Papers:")
for paper in recommended_papers:
st.write(paper)