Update app.py

app.py

@@ -1,83 +1,10 @@
-"https://sites.google.com/airliquide.com/sis-processsafety/knowledge/lessons-learned-sources"
 from sentence_transformers import SentenceTransformer
-
-#model=SentenceTransformer("all-mpnet-base-v2")
+from utils_st import load_models,load_data,clean_whitespace
 import streamlit as st
-import pickle
-import faiss
-from llama_index.core import VectorStoreIndex,StorageContext
-from llama_index.embeddings.huggingface import HuggingFaceEmbedding
-from llama_index.vector_stores.faiss import FaissVectorStore
-from llama_index.core import VectorStoreIndex
-from llama_index.retrievers.bm25 import BM25Retriever
-from llama_index.core.postprocessor import SentenceTransformerRerank
-from llama_index.core import QueryBundle
-from llama_index.core.schema import NodeWithScore
-from llama_index.core.retrievers import BaseRetriever
-from transformers import AutoTokenizer, AutoModel
-
-
-@st.cache_resource(show_spinner=False)
-def load_data():   
-    with open('nodes_clean.pkl', 'rb') as file:
-        nodes=pickle.load( file)
-        d = 768
-        faiss_index = faiss.IndexFlatL2(d)
-        vector_store = FaissVectorStore(faiss_index=faiss_index )
-        storage_context = StorageContext.from_defaults(vector_store=vector_store)
-        # use later nodes_clean
-        index = VectorStoreIndex(nodes,embed_model=embed_model,storage_context=storage_context)
-        retriever_dense = index.as_retriever(similarity_top_k=25,embedding=True)
-        retrieverBM25 = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=25)
-        hybrid_retriever = HybridRetriever(retriever_dense, retrieverBM25)
-        return hybrid_retriever
-
-
-
-@st.cache_resource(show_spinner=False)
-class models():
-    def __init__(self):
-        EMBEDDING_MODEL="BAAI/llm-embedder"
-        self.embed_model = HuggingFaceEmbedding(EMBEDDING_MODEL,device='cpu',)
-        self.reranker = SentenceTransformerRerank(top_n=25, model="BAAI/bge-reranker-base",device='cpu',)
-
-mod=models()
-embed_model=mod.embed_model
-reranker= mod.reranker
-
-class HybridRetriever(BaseRetriever):
-    def __init__(self, vector_retriever, bm25_retriever):
-        self.vector_retriever = vector_retriever
-        self.bm25_retriever = bm25_retriever
-        super().__init__()
-    def _retrieve(self, query, **kwargs):
-        bm25_nodes = self.bm25_retriever.retrieve(query, **kwargs)
-        vector_nodes = self.vector_retriever.retrieve(query, **kwargs)
-        # combine the two lists of nodes
-        all_nodes = []
-        node_ids = set()
-        for n in bm25_nodes + vector_nodes:
-            if n.node.node_id not in node_ids:
-                all_nodes.append(n)
-                node_ids.add(n.node.node_id)
-        return all_nodes
-hybrid_retriever = load_data()
-
-
-
-
-
-import re
-def clean_whitespace(text,k=5):
-    text = text.strip()
-    text=" ".join([i for i in  text.split("\n")[:k] if len(i.strip())>25]+text.split("\n")[k:])
-    text = re.sub(r"\.EU", "", text)
-    #text = re.sub(r"\n+", "\n", text)
-    text = re.sub(r"\s+", " ", text)
-    return text.lower()
-
 
 
+embed_model, reranker = load_models()
+hybrid_retriever = load_data(embed_model)
 
 
 def stream(reranked_nodes,text_size=700):
@@ -92,16 +19,11 @@ def stream(reranked_nodes,text_size=700):
         file_name = i_di[0].metadata['file_name']
         summary = i_di[0].metadata['text']
         url = i_di[0].metadata['doc_url'] 
-        
         st.write(f"**Rank {rank+1}:** {file_name} ")
         st.write(f"- Tittle: [{title}](%s)"% url)
         #st.write("check out this [link](%s)" % url)
         with st.expander(f"Summary"):
             st.write(f"{summary}")
-        #st.write(f"- Summary: {summery}")
-        #st.link_button("Link", url)
-
-        #st.write(f"- URL: {url}")
         with st.expander(f"Extra Text(s) "):
             for n_extra,t in enumerate(i_di[:5]):
                 page_n=t.metadata['page_label'] if "page_label" in t.metadata else 'Unknown'
@@ -110,17 +32,16 @@ def stream(reranked_nodes,text_size=700):
                 st.markdown("""---""")
         st.markdown("""---""")
 
-#stream(reranked_nodes,150)
+
 
 # Function to perform search and return sorted documents
 def perform_search(query):
     if query:
         retrieved_nodes = hybrid_retriever.retrieve(query)
-        reranked_nodes = reranker.postprocess_nodes(
+        reranked_nodes = reranker.predict(
             retrieved_nodes,
-            query_bundle=QueryBundle(
-                query
-            ),)
+            query_bundle=query
+            )
         return reranked_nodes
     else:
         return []
@@ -133,11 +54,9 @@ def main():
     st.title("Information Retrieval System")
     query = st.text_input("Enter your question:")
 
-   
-
     if st.button("Search") or query:
         sorted_docs = perform_search(query)
-        #st.session_state.sorted_docs = sorted_docs
+        st.session_state.sorted_docs = sorted_docs
         
     else:
         sorted_docs = st.session_state.get("sorted_docs", [])
@@ -145,7 +64,7 @@ def main():
 
 
     if sorted_docs:
-        stream(sorted_docs,500)
+        stream(sorted_docs,700)
         #st.write(f"Current Page Number: {page_number}")