import streamlit as st
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from flashrank import Ranker, RerankRequest


@st.cache
def get_embeddings():
    model_name = "BAAI/bge-large-en-v1.5"
    model_kwargs = {'device': 'cpu',"trust_remote_code":True}
    encode_kwargs = {'normalize_embeddings': True} # set True to compute caosine similarity
    model = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs,)
    return model

baai_embeddings = get_embeddings()
kadhal_Server = FAISS.from_local("./",baai_embeddings)
ranker = Ranker(model_name="ms-marco-MiniLM-L-12-v2", cache_dir="/opt")


st.header('kadhalTensor', divider='red')
st.header('_Adhalal Kadhal Seiveer_ is :blue[cool] :sunglasses:')
st.write("Kadhal Engine on Sangam Literature : by Prabakaran Chandran")


with st.form("my_form"):
   st.write("What do want to know about sangam era's love?")
   
   question_input = st.text_input("")

   
   

   # Every form must have a submit button.
   submitted = st.form_submit_button("Submit")
   if submitted:
       
       docs = kadhal_Server.similarity_search(question_input)
       tobeReranked = [{"text":doc.page_content , "metadata":doc.metadata} for doc in docs]
       rerankInput = RerankRequest(
                        passages=tobeReranked,
                        query=" take care of our loved ones accorting to thirukkural?",)
       reranked = ranker.rerank(rerankInput)

       reranked_top = reranked[0:2]
       st.write(reranked_top)

st.write("Outside the form")