Spaces:
Sleeping
Sleeping
Create app.py
Browse files
app.py
ADDED
@@ -0,0 +1,82 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
import streamlit as st
|
2 |
+
import pickle
|
3 |
+
import os
|
4 |
+
import json
|
5 |
+
|
6 |
+
from collections import defaultdict
|
7 |
+
from langchain.vectorstores import FAISS
|
8 |
+
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
|
9 |
+
from rank_bm25 import BM25Okapi
|
10 |
+
|
11 |
+
# Constants
|
12 |
+
BASE_DIR = "built_index"
|
13 |
+
VECTOR_STORE_DIR = os.path.join(BASE_DIR, "vector_store")
|
14 |
+
BM25_INDEX_FILE = os.path.join(BASE_DIR, "bm25_index.pkl")
|
15 |
+
SEARCH_INDEX_FILE = os.path.join(BASE_DIR, "search_index.json")
|
16 |
+
|
17 |
+
# Load embedding model
|
18 |
+
@st.cache_resource
|
19 |
+
def load_embeddings():
|
20 |
+
return HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
|
21 |
+
|
22 |
+
# Load indexes
|
23 |
+
@st.cache_resource
|
24 |
+
def load_indexes():
|
25 |
+
# Load search index
|
26 |
+
with open(SEARCH_INDEX_FILE, "r") as f:
|
27 |
+
index = defaultdict(dict, json.load(f))
|
28 |
+
|
29 |
+
# Load vector store
|
30 |
+
embeddings = load_embeddings()
|
31 |
+
vector_store = FAISS.load_local(VECTOR_STORE_DIR, embeddings, allow_dangerous_deserialization=True)
|
32 |
+
|
33 |
+
# Load BM25 index
|
34 |
+
with open(BM25_INDEX_FILE, "rb") as f:
|
35 |
+
bm25, bm25_texts, url_order = pickle.load(f)
|
36 |
+
|
37 |
+
return index, vector_store, bm25, bm25_texts, url_order
|
38 |
+
|
39 |
+
# Search functions
|
40 |
+
def semantic_search(vector_store, query, k=5):
|
41 |
+
results = vector_store.similarity_search(query, k=k)
|
42 |
+
return [{
|
43 |
+
"url": r.metadata.get("url", "N/A"),
|
44 |
+
"snippet": r.page_content[:200]
|
45 |
+
} for r in results]
|
46 |
+
|
47 |
+
def bm25_search(bm25, bm25_texts, url_order, index, query, k=5):
|
48 |
+
query_tokens = query.lower().split()
|
49 |
+
scores = bm25.get_scores(query_tokens)
|
50 |
+
top_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)[:k]
|
51 |
+
return [{
|
52 |
+
"url": url_order[i],
|
53 |
+
"score": scores[i],
|
54 |
+
"snippet": index[url_order[i]]["content"][:200]
|
55 |
+
} for i in top_indices]
|
56 |
+
|
57 |
+
# Streamlit UI
|
58 |
+
def main():
|
59 |
+
st.set_page_config(page_title="LangChain Search Engine", layout="wide")
|
60 |
+
st.title("🔍 LangChain Search Engine (Semantic + BM25)")
|
61 |
+
|
62 |
+
query = st.text_input("Enter your search query:", "")
|
63 |
+
|
64 |
+
if query:
|
65 |
+
index, vector_store, bm25, bm25_texts, url_order = load_indexes()
|
66 |
+
|
67 |
+
with st.spinner("Searching..."):
|
68 |
+
sem_results = semantic_search(vector_store, query)
|
69 |
+
bm25_results = bm25_search(bm25, bm25_texts, url_order, index, query)
|
70 |
+
|
71 |
+
st.subheader("🔎 Semantic Search Results")
|
72 |
+
for i, res in enumerate(sem_results, 1):
|
73 |
+
st.markdown(f"**{i}. [{res['url']}]({res['url']})**")
|
74 |
+
st.write(res['snippet'] + "...")
|
75 |
+
|
76 |
+
st.subheader("🧮 BM25 Sparse Search Results")
|
77 |
+
for i, res in enumerate(bm25_results, 1):
|
78 |
+
st.markdown(f"**{i}. [{res['url']}]({res['url']})** (Score: {res['score']:.2f})")
|
79 |
+
st.write(res['snippet'] + "...")
|
80 |
+
|
81 |
+
if __name__ == "__main__":
|
82 |
+
main()
|