sss code

app.py

@@ -1,4 +1,301 @@
 import streamlit as st
+import json
+from typing import List
+from fastembed import LateInteractionTextEmbedding, TextEmbedding
+from fastembed import SparseTextEmbedding, SparseEmbedding
+from qdrant_client import QdrantClient, models
+from tokenizers import Tokenizer
 
-x = st.slider('Select a value')
-st.write(x, 'squared is', x * x)
+#############################
+# 1. Utility / Helper Code
+#############################
+
+@st.cache_resource
+def load_tokenizer():
+    """
+    Load the tokenizer for interpreting sparse embeddings (optional usage).
+    """
+    return Tokenizer.from_pretrained(SparseTextEmbedding.list_supported_models()[0]["sources"]["hf"])
+
+@st.cache_resource
+def load_models():
+    """
+    Load/initialize your models once and cache them.
+    """
+    # Dense embedding model
+    dense_embedding_model = TextEmbedding("BAAI/bge-small-en-v1.5")
+
+    # Late interaction model (ColBERTv2)
+    late_embedding_model = LateInteractionTextEmbedding("colbert-ir/colbertv2.0")
+
+    # Sparse embedding model
+    sparse_model_name = "Qdrant/bm25"
+    sparse_model = SparseTextEmbedding(model_name=sparse_model_name)
+
+    return dense_embedding_model, late_embedding_model, sparse_model
+
+def build_qdrant_index(data):
+    """
+    Given the parsed data (list of items), build an in-memory Qdrant index
+    with dense, late, and sparse vectors.
+    """
+    # Extract fields
+    items = data["items"]
+    descriptions = [f"{item['name']} - {item['description']}" for item in items]
+    names = [item["name"] for item in items]
+    metadata = [
+        {"name": item["name"]}  # You can store more fields if you like
+        for item in items
+    ]
+
+    # Load models
+    dense_embedding_model, late_embedding_model, sparse_model = load_models()
+
+    # Generate embeddings
+    dense_embeddings = list(dense_embedding_model.embed(descriptions))
+    name_dense_embeddings = list(dense_embedding_model.embed(names))
+    late_embeddings = list(late_embedding_model.embed(descriptions))
+    sparse_embeddings: List[SparseEmbedding] = list(sparse_model.embed(descriptions, batch_size=6))
+
+    # Create an in-memory Qdrant instance
+    qdrant_client = QdrantClient(":memory:")
+
+    # Create collection schema
+    qdrant_client.create_collection(
+        collection_name="items",
+        vectors_config={
+            "dense": models.VectorParams(
+                size=len(dense_embeddings[0]),
+                distance=models.Distance.COSINE,
+            ),
+            "late": models.VectorParams(
+                size=len(late_embeddings[0][0]),
+                distance=models.Distance.COSINE,
+                multivector_config=models.MultiVectorConfig(
+                    comparator=models.MultiVectorComparator.MAX_SIM
+                ),
+            ),
+        },
+        sparse_vectors_config={
+            "sparse": models.SparseVectorParams(
+                modifier=models.Modifier.IDF,
+            ),
+        }
+    )
+
+    # Upload points
+    points = []
+    for idx, _ in enumerate(metadata):
+        points.append(
+            models.PointStruct(
+                id=idx,
+                payload=metadata[idx],
+                vector={
+                    "late": late_embeddings[idx].tolist(),
+                    "dense": dense_embeddings[idx],
+                    "sparse": sparse_embeddings[idx].as_object(),
+                },
+            )
+        )
+
+    qdrant_client.upload_points(
+        collection_name="items",
+        points=points,
+    )
+
+    return qdrant_client
+
+def run_queries(qdrant_client, query_text):
+    """
+    Run all the different query types and return results in a dictionary.
+    """
+    # Load models
+    dense_embedding_model, late_embedding_model, sparse_model = load_models()
+
+    # Generate single-query embeddings
+    dense_query = next(dense_embedding_model.query_embed(query_text))
+    late_query = next(late_embedding_model.query_embed(query_text))
+    sparse_query = next(sparse_model.query_embed(query_text))
+
+    # For the fusion approach, we need a list form for prefetch
+    tsq = list(sparse_model.embed(query_text, batch_size=6))
+
+    # We'll store top-5 results for each approach
+    results = {}
+
+    # 1) ColBERT (late)
+    results["C"] = qdrant_client.query_points(
+        collection_name="items",
+        query=late_query,
+        using="late",
+        limit=5,
+        with_payload=True
+    )
+
+    # 2) Sparse only
+    results["S"] = qdrant_client.query_points(
+        collection_name="items",
+        query=models.SparseVector(**sparse_query.as_object()),
+        using="sparse",
+        limit=5,
+        with_payload=True
+    )
+
+    # 3) Dense only
+    results["D"] = qdrant_client.query_points(
+        collection_name="items",
+        query=dense_query,
+        using="dense",
+        limit=5,
+        with_payload=True
+    )
+
+    # 4) Hybrid fusion (RRF for Sparse+Dense)
+    results["S+D-F"] = qdrant_client.query_points(
+        collection_name="items",
+        prefetch=[
+            models.Prefetch(
+                query=dense_query,
+                using="dense",
+                limit=100,
+            ),
+            models.Prefetch(
+                query=tsq[0].as_object(),
+                using="sparse",
+                limit=50,
+            )
+        ],
+        query=models.FusionQuery(fusion=models.Fusion.RRF),
+        limit=5,
+        with_payload=True
+    )
+
+    # 5) Hybrid fusion + ColBERT
+    sparse_dense_prefetch = models.Prefetch(
+        prefetch=[
+            models.Prefetch(query=dense_query, using="dense", limit=100),
+            models.Prefetch(query=tsq[0].as_object(), using="sparse", limit=50),
+        ],
+        limit=10,
+        query=models.FusionQuery(fusion=models.Fusion.RRF),
+    )
+    results["S+D-F-C"] = qdrant_client.query_points(
+        collection_name="items",
+        prefetch=[sparse_dense_prefetch],
+        query=late_query,
+        using="late",
+        limit=5,
+        with_payload=True
+    )
+
+    # 6) Hybrid no-fusion + ColBERT
+    old_prefetch = models.Prefetch(
+        prefetch=[
+            models.Prefetch(
+                prefetch=[
+                    models.Prefetch(query=dense_query, using="dense", limit=100)
+                ],
+                query=tsq[0].as_object(),
+                using="sparse",
+                limit=50,
+            )
+        ]
+    )
+    results["S+D-C"] = qdrant_client.query_points(
+        collection_name="items",
+        prefetch=[old_prefetch],
+        query=late_query,
+        using="late",
+        limit=5,
+        with_payload=True
+    )
+
+    return results
+
+#############################
+# 2. Streamlit Main App
+#############################
+
+def main():
+    st.title("Semantic Search Sandbox")
+
+    # Initialize session state if not present
+    if "json_loaded" not in st.session_state:
+        st.session_state["json_loaded"] = False
+    if "qdrant_client" not in st.session_state:
+        st.session_state["qdrant_client"] = None
+
+    #######################################
+    # Show JSON input only if not loaded
+    #######################################
+    if not st.session_state["json_loaded"]:
+        st.subheader("Paste items.json Here")
+        default_json = """
+{
+  "items": [
+    {
+      "name": "Example1",
+      "description": "An example item"
+    },
+    {
+      "name": "Example2",
+      "description": "Another item for demonstration"
+    }
+  ]
+}
+""".strip()
+
+        json_text = st.text_area("JSON Input", value=default_json, height=300)
+
+        if st.button("Load JSON"):
+            try:
+                data = json.loads(json_text)
+                # Build Qdrant index in memory
+                st.session_state["qdrant_client"] = build_qdrant_index(data)
+                st.session_state["json_loaded"] = True
+                st.success("JSON loaded and Qdrant index built successfully!")
+                st.rerun()
+            except Exception as e:
+                st.error(f"Error parsing JSON: {e}")
+
+    else:
+        # The data is loaded, show a button to reset if you want to load new JSON
+        if st.button("Load a different JSON"):
+            st.session_state["json_loaded"] = False
+            st.session_state["qdrant_client"] = None
+            #st.experimental_rerun()  # Refresh the page
+        else:
+            # Show the search interface
+            query_text = st.text_input("Search Query", value="ACB 1.0 Ports")
+            if st.button("Search"):
+                if st.session_state["qdrant_client"] is None:
+                    st.warning("Please load valid JSON first.")
+                    return
+
+                # Run queries
+                results_dict = run_queries(st.session_state["qdrant_client"], query_text)
+
+                # Display results in columns
+                col_names = list(results_dict.keys())
+                # You can split into multiple rows if there are more than 3
+                n_cols = 3
+                # We'll create enough columns to handle all search types
+                rows_needed = (len(col_names) + n_cols - 1) // n_cols
+
+                for row_idx in range(rows_needed):
+                    cols = st.columns(n_cols)
+                    for col_idx in range(n_cols):
+                        method_idx = row_idx * n_cols + col_idx
+                        if method_idx < len(col_names):
+                            method = col_names[method_idx]
+                            qdrant_result = results_dict[method]
+
+                            with cols[col_idx]:
+                                st.markdown(f"### {method}")
+                                for point in qdrant_result.points:
+                                    name = point.payload.get("name", "Unnamed")
+                                    score = round(point.score, 4) if point.score else "N/A"
+                                    st.write(f"- **{name}** (score={score})")
+
+if __name__ == "__main__":
+    main()