Spaces:
Runtime error
Runtime error
File size: 4,283 Bytes
4a02afe 5f39285 4a02afe 5f39285 4a02afe 5f39285 4a02afe | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | """Semantic search with keyword boosting over structured captions."""
from __future__ import annotations
import json
import re
import numpy as np
from sentence_transformers import SentenceTransformer
from caption_store import all_entries
_embed_model: SentenceTransformer | None = None
_MODEL_NAME = "BAAI/bge-base-en-v1.5" # stronger than MiniLM
MIN_RELEVANCE = 0.6
TOP_K = 20
# Weight given to keyword boost relative to semantic score (0–1 additive)
KEYWORD_BOOST = 0.25
def _get_embed_model() -> SentenceTransformer:
global _embed_model
if _embed_model is None:
_embed_model = SentenceTransformer(_MODEL_NAME)
return _embed_model
def _query_tokens(query: str) -> list[str]:
"""Lowercase words from query, 3+ chars."""
return [w for w in re.findall(r"\b\w+\b", query.lower()) if len(w) >= 3]
def _keyword_score(query_tokens: list[str], search_text: str, raw_caption: str) -> float:
"""
Boost score if query tokens appear in high-signal fields (attire, tags, summary).
Returns a value in [0, 1].
"""
if not query_tokens:
return 0.0
# Tokenize the target text strictly by word boundaries to avoid partial substring hits
target_words = set(re.findall(r"\b\w+\b", search_text.lower()))
# Calculate regular hits
hits = sum(1 for token in query_tokens if token in target_words)
# Extra weight for specific high-signal fields (attire, keywords, tags)
high_signal_words = set()
try:
meta = json.loads(raw_caption)
# Gather arrays that explicitly hold high-intent search terms
for field in [meta.get("search_tags", []), meta.get("archive_keywords", []), meta.get("subjects", {}).get("attire", [])]:
if isinstance(field, list):
for item in field:
high_signal_words.update(re.findall(r"\b\w+\b", str(item).lower()))
elif isinstance(field, str):
high_signal_words.update(re.findall(r"\b\w+\b", field.lower()))
except (json.JSONDecodeError, TypeError):
pass
high_signal_hits = sum(1 for token in query_tokens if token in high_signal_words)
# Instead of penalizing long queries, reward ANY valid keyword match heavily
if hits == 0:
return 0.0
base_score = hits / len(query_tokens)
boost_bonus = 0.5 if high_signal_hits > 0 else 0.0
return min(base_score + boost_bonus, 1.0)
def search(query: str, collection: str | None = None) -> list[dict]:
"""
Hybrid search: semantic similarity + keyword boost on structured fields.
Returns up to TOP_K results with score >= MIN_RELEVANCE.
"""
all_indexed = all_entries()
if not all_indexed:
return []
# Filter by collection if specified
if collection and collection != "All":
entries = {
p: data for p, data in all_indexed.items()
if data.get("collection") == collection
}
else:
entries = all_indexed
if not entries:
return []
model = _get_embed_model()
query_tokens = _query_tokens(query)
paths = list(entries.keys())
search_texts = [entries[p].get("search_text") or entries[p]["caption"] for p in paths]
raw_captions = [entries[p]["caption"] for p in paths]
# Query side gets an explicit instruction prefix
query_text = f"Represent this sentence for searching relevant passages: {query}"
query_vec = model.encode([query_text], normalize_embeddings=True)
# Document side stays normal
caption_vecs = model.encode(search_texts, normalize_embeddings=True)
semantic_scores = (caption_vecs @ query_vec.T).flatten()
final_scores = []
for i, (sem, st, rc) in enumerate(zip(semantic_scores, search_texts, raw_captions)):
kw = _keyword_score(query_tokens, st, rc)
final_scores.append(float(sem) + KEYWORD_BOOST * kw)
ranked = sorted(
zip(paths, search_texts, final_scores),
key=lambda x: x[2],
reverse=True,
)
return [
{"path": p, "caption": c, "score": round(s, 4)}
for p, c, s in ranked[:TOP_K]
if s >= MIN_RELEVANCE
]
|