import streamlit as st
import torch
from sentence_transformers import SentenceTransformer, util

embedder = SentenceTransformer('all-mpnet-base-v2')

st.title("iSeBetter: Semantic Issue Transformer")
st.caption("Empower Developers with Advanced AI Issue Resolution")

st.markdown('#')

st.subheader("Search for Similar Issues")
form = st.form("Search for Similar Issues")
text_input = form.text_area("Please enter the issue details",
                            help="Title and description of the issue could be a good start!")

if form.form_submit_button("Run Inference"):
    # Perform semantic search
    query_embedding = embedder.encode(text_input, convert_to_tensor=True)
    corpus_embeddings = torch.load('saved_corpus.pt')
    corpus_embeddings_name = torch.load('saved_corpus_list.txt')
    cos_scores = util.cos_sim(query_embedding, corpus_embeddings)[0]
    top_results = torch.topk(cos_scores, k=5)

    st.markdown('#')
    
    # Display results
    st.subheader("Top 5 Similar Issues")
    for score, idx in zip(top_results[0], top_results[1]):
        st.progress(score.item(), f"Issue: {corpus_embeddings_name[idx]}, Score: {score:.4f}")