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 = model.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)