fix: db optimizations

app.py

@@ -1,10 +1,11 @@
 from huggingface_hub import hf_hub_download
 from gensim.models import Word2Vec
-from nltk import word_tokenize
+from nltk import word_tokenize, sent_tokenize
 from pylatexenc.latex2text import LatexNodes2Text
 
 import faiss
 import duckdb
+import time
 
 import streamlit as st
 import numpy as np
@@ -16,17 +17,35 @@ def get_db(path='arxiv.db'):
     return duckdb.connect(path)
 
 
-def query_neighbours(rows: list):
-    global db
-    con = db
-    rows = [int(x) for x in rows] # Convert numpy.int64 → Python int
-    placeholders = ",".join("?" for _ in rows)
-    df = con.execute(
-        f"SELECT * FROM arxiv WHERE column0 IN ({placeholders})",
-        rows,
-    ).fetchdf()
+@st.cache_resource
+def get_fast_lookup(_model):
+    vectors = _model.wv.vectors  # NumPy matrix (fast)
+    word_to_index = {word: idx for idx, word in enumerate(_model.wv.index_to_key)}
+    return vectors, word_to_index
+
+@st.cache_resource
+def load_arxiv_dict():
+    con = duckdb.connect("arxiv.db")
+    df = con.execute("""
+        SELECT column0, id, title, abstract, categories
+        FROM arxiv
+    """).fetchdf()
+
+    # dictionary: column0 → row
+    return {
+        int(row["column0"]): {
+            "id": row["id"],
+            "title": row["title"],
+            "abstract": row["abstract"],
+            "categories": row["categories"]
+        }
+        for _, row in df.iterrows()
+    }
+
+def query_neighbours(rows):
+    global arxiv_dict
+    return [arxiv_dict.get(int(x)) for x in rows if int(x) in arxiv_dict]
 
-    return df.to_dict("records")
 
 @st.cache_resource
 def get_model():
@@ -52,24 +71,34 @@ def get_faiss_index():
 
 
 def run_semantic_search(query, top_k):
-    global model
-    global faiss_index
+    global model, faiss_index, word_to_index, vectors
 
     index = faiss_index
 
-    words = word_tokenize(query.lower())
+    words = query.lower().split()
     vecs = []
 
+    start_t = time.time()
     for w in words:
-        if w in model.wv:
-            vecs.append(model.wv[w])
-    if len(vecs) == 0:
+        idx = word_to_index.get(w)
+        if idx is not None:
+            vecs.append(vectors[idx])
+    mid_t = time.time()
+    print(f"Tokenization time: {mid_t - start_t}")
+
+    if not vecs:
         return []
+
     qvec = np.mean(vecs, axis=0).astype('float32').reshape(1, -1)
     faiss.normalize_L2(qvec)
+
     scores, neighbors = index.search(qvec, top_k)
+    mid2_t = time.time()
+    print(f"Search time : {mid2_t - mid_t}")
+    result = query_neighbours(neighbors[0])
+    print(f"Query time : {time.time() - mid2_t}\n\n\n")
+    return result
 
-    return query_neighbours(neighbors[0])
 
 
 #-----------------------------------
@@ -79,6 +108,8 @@ def run_semantic_search(query, top_k):
 model = get_model()
 faiss_index = get_faiss_index()
 db = get_db()
+vectors, word_to_index = get_fast_lookup(model)
+arxiv_dict = load_arxiv_dict()
 
 # ----------------------------------
 # Streamlit Page Setup
@@ -105,20 +136,24 @@ search_button = st.button("Search")
 # Handle search click
 # --------------------------------------------------------------
 if search_button and query.strip():
+    start_time = time.time()
     with st.spinner("Searching..."):
         results = run_semantic_search(query, top_k)
-
-    st.header(f"Top {top_k} Results")
-
-    # ----------------------------------------------------------
-    # Display results (card-style)
-    # ----------------------------------------------------------
-    for i, paper in enumerate(results, start=1):
-        st.markdown(f"### **{i}. {LatexNodes2Text().latex_to_text(paper['title'].replace("\n", " ").strip())}**")
-
-        st.markdown(f"**Categories:** {paper['categories']}")
-        st.markdown(f"**Abstract:** {LatexNodes2Text().latex_to_text(paper["abstract"][:600])}...")
-        st.markdown(f"[View on arXiv](https://arxiv.org/abs/{paper['id']})")
-
-        st.markdown("---")
+    end_time = time.time()
+    elapsed = end_time - start_time
+    st.write(f"**Your query took {elapsed:.3f} seconds**")
+    if(len(results) != 0):
+        st.header(f"Top {top_k} Results")
+
+        # ----------------------------------------------------------
+        # Display results (card-style)
+        # ----------------------------------------------------------
+        for i, paper in enumerate(results, start=1):
+            st.markdown(f"### **[{i}. {LatexNodes2Text().latex_to_text(paper['title'].replace("\n", " ").strip())}](https://arxiv.org/abs/{paper['id']})**")
+
+            st.markdown(f"**Categories:** {paper['categories']}")
+            st.markdown(f"**Abstract:** {LatexNodes2Text().latex_to_text(paper["abstract"][:600])}...")
+            st.markdown("---")
+    else:
+        st.markdown(f"No Results, either model is not trained on this word")