Upload 2 files

Test_SemanticSearch.py

@@ -0,0 +1,149 @@
+import numpy as np
+import pandas as pd
+from pathlib import Path
+import streamlit as st
+import pickle
+import plotly.express as px
+import plotly.graph_objects as go 
+import gdown
+import torch
+from sentence_transformers import SentenceTransformer, util
+
+# -------------------------------
+# Setup
+# -------------------------------
+st.set_page_config(page_title="Alibaba Semantic Search", layout="wide")
+
+MODEL_DIR = Path("models")
+MODEL_DIR.mkdir(exist_ok=True)
+
+embeddings_path = MODEL_DIR / 'desc_embeddings_Alibaba_20251016.npy' 
+umap_embeddings_path = MODEL_DIR / 'descs_umap_2d_AB_20251016.npy' 
+data_file_path = MODEL_DIR / 'full_df_minus_nan_descs.csv' 
+umap_model_path = MODEL_DIR / 'umap_2d_AB_written.pkl' 
+pca_model_path = MODEL_DIR / 'pca_AB_written.pkl' 
+
+emb_ID = '1QQ_QfFTSzTLNkp6Sr4jux_ZTJjMhSyah'
+umap_emb_ID = '1a5t5iWOAVgUmYXzrWXctATkDyx9rRF4F'
+data_ID = '1tzM67Lg3R-rAvRtol0VGHx6zGW_tdx60'
+umap_mod_ID = '1x8PK1Gn72YYBZ4po-0guZMUBtL8oSn1i'
+pca_mod_ID = '1jIxBBAZOy8OAzGxBCG4jy7244Wb_TjP9'
+
+paths = [embeddings_path, umap_embeddings_path, data_file_path, umap_model_path, pca_model_path]
+ids = [emb_ID, umap_emb_ID, data_ID, umap_mod_ID, pca_mod_ID]
+assets_links = [f"https://drive.google.com/uc?id={x}" for x in ids]
+
+# -------------------------------
+# Download + Load Data
+# -------------------------------
+def load_assets():
+    st.info("Downloading assets from Google Drive (only if missing)...")
+    for url, path in zip(assets_links, paths):
+        if not path.exists():
+            gdown.download(url, str(path), quiet=False)
+    st.success("Assets ready ✅")
+
+    embeddings = np.load(embeddings_path)
+    umap_2d = np.load(umap_embeddings_path)
+    docs = pd.read_csv(data_file_path)
+    with open(umap_model_path, "rb") as f:
+        umap_model = pickle.load(f)
+    with open(pca_model_path, "rb") as f:
+        pca_model = pickle.load(f) 
+    return embeddings, umap_2d, docs, umap_model, pca_model
+
+embeddings, umap_2d, docs, umap_model, pca_model = load_assets()
+
+# -------------------------------
+# Load SentenceTransformer (cached)
+# -------------------------------
+@st.cache_resource
+def load_text_encoder():
+    return SentenceTransformer('Alibaba-NLP/gte-multilingual-base', trust_remote_code=True)
+
+model = load_text_encoder()
+
+# -------------------------------
+# UI
+# -------------------------------
+st.title("🔍 Semantic Search — Alibaba Embeddings")
+st.markdown("Enter a query to highlight semantically similar documents on the 2D UMAP plot.")
+
+query = st.text_input("Enter search query:")
+top_k = st.slider("Number of matches to highlight", min_value=10, max_value=2500, value=100)
+
+similarity_measure = st.radio(
+    "Similarity measure",
+    ["Cosine", "Euclidean", "Manhattan (L1)"],
+    horizontal=True
+)
+
+# -------------------------------
+# Search logic
+# -------------------------------
+if query:
+    with st.spinner("Encoding and searching..."):
+        query_embedding = model.encode(query, convert_to_tensor=True)
+        query_numpy = query_embedding.cpu().numpy().reshape(1, -1)
+        query_pca = pca_model.transform(query_numpy)
+        query_umap = umap_model.transform(query_pca)
+
+        if similarity_measure == "Cosine":
+            scores = util.cos_sim(query_embedding, embeddings)[0]
+        elif similarity_measure == "Euclidean":
+            scores = -torch.cdist(query_embedding, embeddings, p=2)[0]
+        elif similarity_measure == "Manhattan (L1)":
+            scores = -torch.cdist(query_embedding, embeddings, p=1)[0]
+
+        top_results = scores.argsort(descending=True)
+        highlight_indices = top_results[:top_k].cpu().numpy()
+
+    documents = docs.title_narrative
+    reporting_org = docs.reporting_org_name
+    funding = docs.Funding
+
+    labels = ["Match" if i in highlight_indices else "Other" for i in range(len(documents))]
+
+    df = pd.DataFrame({
+        "UMAP_1": umap_2d[:, 0],
+        "UMAP_2": umap_2d[:, 1],
+        "Label": labels,
+        "Text": documents,
+        "Reporting Org": reporting_org,
+        "Funding": funding,
+    })
+
+    df["Title"] = df["Text"].str.slice(0, 100) + "..."
+    df["Index"] = df.index
+
+    color_discrete_map = {"Match": "red", "Other": "lightgray"}
+
+    fig = px.scatter(
+        df,
+        x="UMAP_1",
+        y="UMAP_2",
+        color="Label",
+        color_discrete_map=color_discrete_map,
+        hover_data={"Text": False, "Title": True, "Index": True, "Reporting Org": True, "Funding":True, "UMAP_1": False, "UMAP_2": False},
+        opacity=0.7,
+        title=f"Top {top_k} semantic matches for: '{query}' ({similarity_measure})",
+        width=900,
+        height=700
+    )
+
+    fig.add_trace(go.Scatter(
+        x=[query_umap[0][0]], y=[query_umap[0][1]],
+        mode='markers+text',
+        marker=dict(size=10, color='blue', symbol='x'),
+        name='Query',
+        text=['Query'], textposition='top center'
+    ))
+
+    fig.update_traces(marker=dict(size=4))
+    st.plotly_chart(fig, use_container_width=True)
+
+    st.subheader("Top 10 matched documents")
+    for rank, idx in enumerate(highlight_indices[:10], start=1):
+        st.markdown(f"{rank}. {documents.iloc[idx]}")
+else:
+    st.info("Enter a search query to begin.")

requirements.txt

@@ -1,3 +1,15 @@
-altair
-pandas
-streamlit
+# Core app
+streamlit==1.39.0
+gdown==5.2.0
+pandas==2.2.3
+numpy==2.1.2
+numba==0.61.2
+joblib==1.4.2
+plotly==5.24.1
+torch==2.6.0
+sentence-transformers==3.2.1
+umap-learn==0.5.6
+scikit-learn==1.5.2
+
+# Optional (used internally by SentenceTransformer models)
+transformers==4.45.2