Spaces:
Running on Zero
Running on Zero
Deploy from sec-dense-fts
Browse files- README.md +161 -4
- app.py +516 -0
- fts_queries.py +386 -0
- requirements.txt +5 -0
- search_core.py +237 -0
README.md
CHANGED
|
@@ -1,13 +1,170 @@
|
|
| 1 |
---
|
| 2 |
title: SEC FTS Hybrid Comparison
|
| 3 |
-
emoji:
|
| 4 |
-
colorFrom:
|
| 5 |
-
colorTo:
|
| 6 |
sdk: gradio
|
| 7 |
sdk_version: 6.29.0
|
| 8 |
python_version: '3.12'
|
| 9 |
app_file: app.py
|
| 10 |
pinned: false
|
|
|
|
| 11 |
---
|
| 12 |
|
| 13 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
---
|
| 2 |
title: SEC FTS Hybrid Comparison
|
| 3 |
+
emoji: ๐
|
| 4 |
+
colorFrom: blue
|
| 5 |
+
colorTo: indigo
|
| 6 |
sdk: gradio
|
| 7 |
sdk_version: 6.29.0
|
| 8 |
python_version: '3.12'
|
| 9 |
app_file: app.py
|
| 10 |
pinned: false
|
| 11 |
+
short_description: Pinecone full-text search vs dense vectors on 10-Ks
|
| 12 |
---
|
| 13 |
|
| 14 |
+
# SEC Document Search โ Pinecone FTS + Semantic
|
| 15 |
+
|
| 16 |
+
A demo application showing how to build a hybrid search system over SEC 10-K filings using [Pinecone's full-text search](https://docs.pinecone.io/guides/search/full-text-search) combined with dense vector embeddings.
|
| 17 |
+
|
| 18 |
+
The app supports three search modes:
|
| 19 |
+
|
| 20 |
+
| Mode | How it works |
|
| 21 |
+
|---|---|
|
| 22 |
+
| **Full-text** | BM25 keyword search over the document text |
|
| 23 |
+
| **Query builder (full-text)** | Construct complex full-text queries โ Lucene boolean / phrase / proximity / boost / prefix scoring plus text-match and metadata filters โ starting from 10 worked examples |
|
| 24 |
+
| **Semantic** | Dense vector similarity via OpenAI embeddings |
|
| 25 |
+
| **Compare: dense vs full-text** | Runs a dense search and a full-text search side by side and contrasts them โ overlap, rank shifts, keyword coverage, latency, and an RRF-fused list |
|
| 26 |
+
| **Hybrid** | Dense vector ranking with a must-contain keyword filter โ results are semantically ranked and guaranteed to contain the specified terms |
|
| 27 |
+
|
| 28 |
+
## Data
|
| 29 |
+
|
| 30 |
+
`example_data/` contains chunked 10-K filings for six companies across six years:
|
| 31 |
+
|
| 32 |
+
| Ticker | Company | Years |
|
| 33 |
+
|---|---|---|
|
| 34 |
+
| AAPL | Apple | 2019โ2024 |
|
| 35 |
+
| AMZN | Amazon | 2019โ2024 |
|
| 36 |
+
| F | Ford | 2019โ2024 |
|
| 37 |
+
| GM | General Motors | 2019โ2024 |
|
| 38 |
+
| MSFT | Microsoft | 2019โ2024 |
|
| 39 |
+
| ORCL | Oracle | 2019โ2024 |
|
| 40 |
+
|
| 41 |
+
~24,000 document chunks total. Each chunk has an `_id`, `text`, `ticker`, `filing_type`, `year`, and `chunk_index`.
|
| 42 |
+
|
| 43 |
+
## Prerequisites
|
| 44 |
+
|
| 45 |
+
- Python 3.12+
|
| 46 |
+
- [`uv`](https://docs.astral.sh/uv/)
|
| 47 |
+
- A [Pinecone API key](https://app.pinecone.io/) (free tier works)
|
| 48 |
+
- An [OpenAI API key](https://platform.openai.com/api-keys)
|
| 49 |
+
|
| 50 |
+
## Setup
|
| 51 |
+
|
| 52 |
+
```bash
|
| 53 |
+
git clone https://github.com/tim-pinecone/sec-dense-fts
|
| 54 |
+
cd sec-dense-fts
|
| 55 |
+
|
| 56 |
+
cp .env.example .env
|
| 57 |
+
# Add your PINECONE_API_KEY and OPENAI_API_KEY to .env
|
| 58 |
+
|
| 59 |
+
uv sync
|
| 60 |
+
```
|
| 61 |
+
|
| 62 |
+
## Ingest
|
| 63 |
+
|
| 64 |
+
Creates the Pinecone index, embeds all chunks with `text-embedding-3-small`, and upserts them in batches. Safe to re-run โ skips index creation if it already exists.
|
| 65 |
+
|
| 66 |
+
```bash
|
| 67 |
+
uv run main.py
|
| 68 |
+
```
|
| 69 |
+
|
| 70 |
+
Ingestion takes a few minutes (OpenAI embedding calls are the bottleneck). The script polls until all documents are searchable before exiting.
|
| 71 |
+
|
| 72 |
+
## Run the app
|
| 73 |
+
|
| 74 |
+
Two UIs share the same search logic (`search_core.py`, `fts_queries.py`):
|
| 75 |
+
|
| 76 |
+
| UI | Command | URL | Notes |
|
| 77 |
+
|---|---|---|---|
|
| 78 |
+
| Gradio | `uv run python app.py` | [http://localhost:7860](http://localhost:7860) | What the Hugging Face Space runs |
|
| 79 |
+
| Streamlit | `uv run streamlit run app_streamlit.py` | [http://localhost:8501](http://localhost:8501) | Local development UI |
|
| 80 |
+
|
| 81 |
+
Both have the same five tabs/modes. Use the sidebar to filter by ticker and year, pick a search mode, and enter your query โ or pick one of the prepared **Example queries** at the top of each mode (three per mode; ten in the query builder). Examples live in `fts_queries.py` (`MODE_EXAMPLES`, `EXAMPLES`).
|
| 82 |
+
|
| 83 |
+
## Full-text query builder
|
| 84 |
+
|
| 85 |
+
The **Query builder** mode exposes the full Pinecone FTS query surface. Queries are built in three parts, and the exact `documents.search(...)` request is shown as JSON and Python before running:
|
| 86 |
+
|
| 87 |
+
1. **Scoring** โ either BM25 keywords (`text`) or Lucene (`query_string`). Lucene can be written raw or assembled row-by-row with the clause builder (MUST `+` / MUST NOT `-` / SHOULD, term / phrase / phrase prefix, slop `~N`, boost `^N`).
|
| 88 |
+
2. **Text-match filters** โ `$match_phrase`, `$match_all`, `$match_any`, each optionally negated with `$not`, combined with `$and` or `$or`.
|
| 89 |
+
3. **Metadata filters** โ ticker `$in` / `$nin`, year range, chunk-index range.
|
| 90 |
+
|
| 91 |
+
Pick an example from the dropdown to load it into the builder, then tweak it:
|
| 92 |
+
|
| 93 |
+
| Example | Demonstrates |
|
| 94 |
+
|---|---|
|
| 95 |
+
| Supply-chain shortages, excluding COVID | `text:(("supply chain" OR semiconductor) AND shortage) NOT text:(covid)` |
|
| 96 |
+
| Cyber incidents | `text:(+cybersecurity ransomware^3 breach -insurance)` |
|
| 97 |
+
| Rising interest rates | Proximity: `text:("interest rates increase"~5)` |
|
| 98 |
+
| China trade & tariffs | Term boost: `text:(tariffs^3 trade china)` |
|
| 99 |
+
| AI mentions | Phrase prefix `"artificial intel"*` + required term + year range |
|
| 100 |
+
| EV batteries at Ford & GM | BM25 + `$match_phrase` filter + ticker `$in` + year range |
|
| 101 |
+
| Cloud growth, no pandemic talk | `$not` + `$match_any` exclusion, ticker `$nin` |
|
| 102 |
+
| Regulators: EC or DOJ | `$or` across two `$match_phrase` filters |
|
| 103 |
+
| Inflation & input costs | `$match_all` + year range |
|
| 104 |
+
| Buybacks vs. dividends | Required OR-groups `+(a OR b) +(c OR d)` |
|
| 105 |
+
|
| 106 |
+
Things the server enforces (surfaced in the UI):
|
| 107 |
+
|
| 108 |
+
- One scoring type per request โ `text` **or** `query_string`, never mixed.
|
| 109 |
+
- `query_string` clauses may not set `fields`; qualify terms inline (`text:(...)`).
|
| 110 |
+
- A Lucene query of only exclusions (`text:(-covid)`) is rejected.
|
| 111 |
+
- Proximity, boost and phrase prefix are scoring-only โ they can't be used in `filter`.
|
| 112 |
+
- Phrase-prefix matches all receive the same constant score.
|
| 113 |
+
|
| 114 |
+
The query compilation logic and examples live in `fts_queries.py`.
|
| 115 |
+
|
| 116 |
+
## Comparing dense vs full-text
|
| 117 |
+
|
| 118 |
+
**Compare** mode runs both searches in parallel for the same question:
|
| 119 |
+
|
| 120 |
+
- **Dense side** โ the query is embedded with `text-embedding-3-small` and ranked by cosine similarity.
|
| 121 |
+
- **Full-text side** โ either the same query as BM25 keywords, or whatever complex query is currently set up in the **Query builder** (Lucene, text-match filters, metadata filters). When using the builder, you can optionally apply its filters to the dense side too, so only the ranking signal differs.
|
| 122 |
+
|
| 123 |
+
Sidebar ticker/year filters apply to both sides. The results show:
|
| 124 |
+
|
| 125 |
+
| View | What it tells you |
|
| 126 |
+
|---|---|
|
| 127 |
+
| Overlap / Jaccard, only-dense, only-full-text | How much the two retrieval methods agree in the top-k |
|
| 128 |
+
| Latency | Dense (embedding + search) vs full-text search time |
|
| 129 |
+
| Keyword coverage | Share of the query's keywords present in each result โ dense results with low coverage are paraphrase / concept matches that BM25 can't find |
|
| 130 |
+
| Side by side | Both ranked lists, with badges showing each result's rank in the other list, and keyword highlighting on both |
|
| 131 |
+
| Rank comparison | One table of every retrieved chunk with its dense rank, full-text rank and ฮ |
|
| 132 |
+
| Fused (RRF) | Client-side reciprocal rank fusion (k = 60) of the two lists โ a preview of a two-query hybrid |
|
| 133 |
+
|
| 134 |
+
## Deploying to Hugging Face Spaces
|
| 135 |
+
|
| 136 |
+
The Space runs the **Gradio** app (`app.py`) on free ZeroGPU hardware โ no Docker needed. The YAML block at the top of this README is the Space config, and `requirements.txt` holds the Space's Python dependencies (Gradio itself comes from `sdk_version`). ZeroGPU requires at least one `@spaces.GPU` function; `app.py` defines a no-op one since all compute happens in Pinecone and OpenAI.
|
| 137 |
+
|
| 138 |
+
1. Add a Hugging Face write token to `.env` as `HF_TOKEN`.
|
| 139 |
+
2. Preview what will be uploaded:
|
| 140 |
+
```bash
|
| 141 |
+
uv run python deploy_space.py <owner>/<space-name> --dry-run
|
| 142 |
+
```
|
| 143 |
+
3. Deploy. The first run creates the Space; `--set-secrets` copies `PINECONE_API_KEY` and `OPENAI_API_KEY` from `.env` into the Space's secrets (only needed once, or when keys change):
|
| 144 |
+
```bash
|
| 145 |
+
uv run python deploy_space.py <owner>/<space-name> --set-secrets [--private]
|
| 146 |
+
```
|
| 147 |
+
|
| 148 |
+
`deploy_space.py` uploads an explicit allowlist (`README.md`, `requirements.txt`, `app.py`, `search_core.py`, `fts_queries.py`), so `.env`, the example data, and the Streamlit app never leave your machine. Re-run it without `--set-secrets` to push code changes.
|
| 149 |
+
|
| 150 |
+
## Index schema
|
| 151 |
+
|
| 152 |
+
```
|
| 153 |
+
text โ full-text search field (BM25, English, stemming enabled)
|
| 154 |
+
embedding โ dense vector, 1536 dims, cosine similarity
|
| 155 |
+
ticker โ filterable metadata (string)
|
| 156 |
+
filing_type โ filterable metadata (string)
|
| 157 |
+
year โ filterable metadata (integer)
|
| 158 |
+
chunk_index โ filterable metadata (integer)
|
| 159 |
+
```
|
| 160 |
+
|
| 161 |
+
The FTS and vector fields are declared in the schema at index creation. Metadata fields (`ticker`, `filing_type`, `year`, `chunk_index`) are automatically indexed โ they do not need to be declared.
|
| 162 |
+
|
| 163 |
+
## How hybrid search works
|
| 164 |
+
|
| 165 |
+
The hybrid mode uses a single Pinecone query:
|
| 166 |
+
|
| 167 |
+
- `score_by` โ dense vector cosine similarity (semantic ranking)
|
| 168 |
+
- `filter` โ `$match_all` on the text field (hard lexical requirement)
|
| 169 |
+
|
| 170 |
+
This means results are ordered by semantic relevance, but only chunks that contain all the specified keywords are returned. It's useful for queries like "what does MSFT say about Azure capital expenditure" โ the semantic query captures the intent, and the text filter ensures the specific terms are present.
|
app.py
ADDED
|
@@ -0,0 +1,516 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import html
|
| 2 |
+
import json
|
| 3 |
+
import os
|
| 4 |
+
|
| 5 |
+
import gradio as gr
|
| 6 |
+
import pandas as pd
|
| 7 |
+
|
| 8 |
+
import fts_queries as fq
|
| 9 |
+
from search_core import (
|
| 10 |
+
INCLUDE_FIELDS,
|
| 11 |
+
RRF_K,
|
| 12 |
+
TICKERS,
|
| 13 |
+
YEARS,
|
| 14 |
+
compare_stats,
|
| 15 |
+
highlight,
|
| 16 |
+
keyword_coverage,
|
| 17 |
+
match_record,
|
| 18 |
+
metadata_filter,
|
| 19 |
+
missing_env,
|
| 20 |
+
rank_badge,
|
| 21 |
+
rank_rows,
|
| 22 |
+
rrf_fuse,
|
| 23 |
+
run_compare,
|
| 24 |
+
run_search_hybrid,
|
| 25 |
+
run_search_semantic,
|
| 26 |
+
run_search_text,
|
| 27 |
+
search,
|
| 28 |
+
)
|
| 29 |
+
|
| 30 |
+
os.environ["GRADIO_SSR_MODE"] = "False"
|
| 31 |
+
|
| 32 |
+
try:
|
| 33 |
+
import spaces
|
| 34 |
+
|
| 35 |
+
@spaces.GPU
|
| 36 |
+
def zero_gpu_noop():
|
| 37 |
+
return True
|
| 38 |
+
except Exception:
|
| 39 |
+
def zero_gpu_noop():
|
| 40 |
+
return False
|
| 41 |
+
|
| 42 |
+
|
| 43 |
+
BLANK = "โ Start from scratch โ"
|
| 44 |
+
N_CLAUSES = 6
|
| 45 |
+
N_TEXT_FILTERS = 4
|
| 46 |
+
ANY_YEAR = "Any"
|
| 47 |
+
BUILDER_EXAMPLES = {ex["name"]: ex for ex in fq.EXAMPLES}
|
| 48 |
+
|
| 49 |
+
CSS = """
|
| 50 |
+
.card {border: 1px solid var(--border-color-primary); border-radius: 8px; padding: 10px 12px;
|
| 51 |
+
margin-bottom: 8px; background: var(--block-background-fill);}
|
| 52 |
+
.card .hdr {display: flex; justify-content: space-between; gap: 8px; font-size: 14px;}
|
| 53 |
+
.card .score {font-variant-numeric: tabular-nums; color: var(--body-text-color-subdued);}
|
| 54 |
+
.card .badge {font-size: 12px; color: var(--body-text-color-subdued); margin-top: 2px;}
|
| 55 |
+
.card .snip {font-size: 13px; line-height: 1.45; margin-top: 6px;}
|
| 56 |
+
.card details {font-size: 13px; margin-top: 4px;}
|
| 57 |
+
.card mark, .snip mark {background: rgba(250, 204, 21, .45); color: inherit; padding: 0 1px; border-radius: 2px;}
|
| 58 |
+
.count {font-size: 13px; color: var(--body-text-color-subdued); margin: 4px 0 8px;}
|
| 59 |
+
.err {border: 1px solid #dc2626; border-radius: 8px; padding: 10px 12px; color: #dc2626;}
|
| 60 |
+
.metrics {display: grid; grid-template-columns: repeat(auto-fit, minmax(150px, 1fr)); gap: 8px; margin: 4px 0 8px;}
|
| 61 |
+
.metric {border: 1px solid var(--border-color-primary); border-radius: 8px; padding: 8px 10px;}
|
| 62 |
+
.metric .label {font-size: 12px; color: var(--body-text-color-subdued);}
|
| 63 |
+
.metric .value {font-size: 22px; font-weight: 600; font-variant-numeric: tabular-nums;}
|
| 64 |
+
.metric .sub {font-size: 12px; color: var(--body-text-color-subdued);}
|
| 65 |
+
.cols {display: grid; grid-template-columns: 1fr 1fr; gap: 12px;}
|
| 66 |
+
@media (max-width: 800px) {.cols {grid-template-columns: 1fr;}}
|
| 67 |
+
"""
|
| 68 |
+
|
| 69 |
+
CHEAT_SHEET = """
|
| 70 |
+
| Goal | Lucene (`query_string` scoring) | Filter operator (hard constraint, no scoring) |
|
| 71 |
+
|---|---|---|
|
| 72 |
+
| Any of these terms | `text:(cloud revenue)` | `{"text": {"$match_any": "cloud revenue"}}` |
|
| 73 |
+
| All terms, any order | `text:(+cloud +revenue)` | `{"text": {"$match_all": "cloud revenue"}}` |
|
| 74 |
+
| Exact phrase | `text:("supply chain")` | `{"text": {"$match_phrase": "supply chain"}}` |
|
| 75 |
+
| Exclude | `text:(cloud -covid)` / `NOT text:(covid)` | `{"$not": {"text": {"$match_any": "covid"}}}` |
|
| 76 |
+
| Boolean nesting | `text:((a OR b) AND c)` | `{"$or": [...]}`, `{"$and": [...]}` |
|
| 77 |
+
| Proximity (within N words) | `text:("interest rates increase"~5)` | โ scoring only |
|
| 78 |
+
| Boost a term | `text:(tariffs^3 trade)` | โ scoring only |
|
| 79 |
+
| Phrase prefix | `text:("artificial intel"*)` (โฅ 2 words) | โ scoring only |
|
| 80 |
+
|
| 81 |
+
- One scoring type per request: BM25 `text` **or** `query_string`, never both.
|
| 82 |
+
- Filters run first and shrink the candidate set; scoring ranks what's left.
|
| 83 |
+
- The `text` field uses English stemming, so `increase` also matches `increases` / `increased`.
|
| 84 |
+
- A Lucene query made only of exclusions (`-covid`) is rejected โ include at least one positive clause.
|
| 85 |
+
- Phrase-prefix matches all get the same constant score.
|
| 86 |
+
"""
|
| 87 |
+
|
| 88 |
+
|
| 89 |
+
def card_html(r: dict, terms: list[str], rank: int | None = None, badge: str | None = None) -> str:
|
| 90 |
+
snippet, full = highlight(r["text"], terms)
|
| 91 |
+
prefix = f"<b>#{rank}</b> ยท " if rank else ""
|
| 92 |
+
meta = f"{prefix}<b>{html.escape(r['ticker'])}</b> ยท {r['year']} ยท {html.escape(r['filing'])} ยท chunk {r['chunk']}"
|
| 93 |
+
badge_html = f"<div class='badge'>{html.escape(badge)}</div>" if badge else ""
|
| 94 |
+
more = f"<details><summary>Full text</summary>{full}</details>" if len(r["text"]) > 300 else ""
|
| 95 |
+
return (
|
| 96 |
+
f"<div class='card'><div class='hdr'><span>{meta}</span><span class='score'>{r['score']:.4f}</span></div>"
|
| 97 |
+
f"{badge_html}<div class='snip'>{snippet}</div>{more}</div>"
|
| 98 |
+
)
|
| 99 |
+
|
| 100 |
+
|
| 101 |
+
def results_html(records: list[dict], terms: list[str]) -> str:
|
| 102 |
+
if not records:
|
| 103 |
+
return "<div class='count'>No results found.</div>"
|
| 104 |
+
cards = "".join(card_html(r, terms) for r in records)
|
| 105 |
+
return f"<div class='count'>{len(records)} result(s)</div>{cards}"
|
| 106 |
+
|
| 107 |
+
|
| 108 |
+
def error_html(message: str) -> str:
|
| 109 |
+
return f"<div class='err'>{html.escape(message)}</div>"
|
| 110 |
+
|
| 111 |
+
|
| 112 |
+
def metric_html(label: str, value: str, sub: str = "") -> str:
|
| 113 |
+
sub_html = f"<div class='sub'>{html.escape(sub)}</div>" if sub else ""
|
| 114 |
+
return f"<div class='metric'><div class='label'>{label}</div><div class='value'>{value}</div>{sub_html}</div>"
|
| 115 |
+
|
| 116 |
+
|
| 117 |
+
def mode_example(mode: str, name: str, fields: list[str]):
|
| 118 |
+
ex = next((e for e in fq.MODE_EXAMPLES[mode] if e["name"] == name), None)
|
| 119 |
+
if not ex:
|
| 120 |
+
return [gr.update() for _ in fields] + [""]
|
| 121 |
+
return [ex["values"].get(f, "") for f in fields] + [f"โน๏ธ {ex['description']}"]
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
def do_fulltext(query, tickers, years, top_k):
|
| 125 |
+
if not query.strip():
|
| 126 |
+
return error_html("Enter a keyword query.")
|
| 127 |
+
try:
|
| 128 |
+
resp = run_search_text(query, tickers, years, int(top_k))
|
| 129 |
+
except Exception as e:
|
| 130 |
+
return error_html(f"Search failed: {e}")
|
| 131 |
+
terms = fq.highlight_terms({"scoring": "text", "bm25": query})
|
| 132 |
+
return results_html([match_record(m) for m in resp.matches], terms)
|
| 133 |
+
|
| 134 |
+
|
| 135 |
+
def do_semantic(query, tickers, years, top_k):
|
| 136 |
+
if not query.strip():
|
| 137 |
+
return error_html("Enter a semantic query.")
|
| 138 |
+
try:
|
| 139 |
+
resp = run_search_semantic(query, tickers, years, int(top_k))
|
| 140 |
+
except Exception as e:
|
| 141 |
+
return error_html(f"Search failed: {e}")
|
| 142 |
+
return results_html([match_record(m) for m in resp.matches], [])
|
| 143 |
+
|
| 144 |
+
|
| 145 |
+
def do_hybrid(sem_query, txt_filter, tickers, years, top_k):
|
| 146 |
+
if not sem_query.strip():
|
| 147 |
+
return error_html("Enter a semantic query.")
|
| 148 |
+
try:
|
| 149 |
+
resp = run_search_hybrid(sem_query, txt_filter, tickers, years, int(top_k))
|
| 150 |
+
except Exception as e:
|
| 151 |
+
return error_html(f"Search failed: {e}")
|
| 152 |
+
terms = fq.highlight_terms({"scoring": "text", "bm25": txt_filter})
|
| 153 |
+
return results_html([match_record(m) for m in resp.matches], terms)
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def builder_state(values: list) -> dict:
|
| 157 |
+
it = iter(values)
|
| 158 |
+
scoring, bm25, lucene_mode, raw = next(it), next(it), next(it), next(it)
|
| 159 |
+
clauses = []
|
| 160 |
+
for _ in range(N_CLAUSES):
|
| 161 |
+
occur, kind, value, slop, boost = (next(it) for _ in range(5))
|
| 162 |
+
if str(value or "").strip():
|
| 163 |
+
clauses.append({"occur": occur, "kind": kind, "value": value, "slop": slop or 0, "boost": boost or 1.0})
|
| 164 |
+
text_filters = []
|
| 165 |
+
for _ in range(N_TEXT_FILTERS):
|
| 166 |
+
op, value, negate = next(it), next(it), next(it)
|
| 167 |
+
if str(value or "").strip():
|
| 168 |
+
text_filters.append({"op": op, "value": value, "negate": bool(negate)})
|
| 169 |
+
join, tickers, excl, year_from, year_to, chunk_min, chunk_max = (next(it) for _ in range(7))
|
| 170 |
+
|
| 171 |
+
year_range = None
|
| 172 |
+
if year_from != ANY_YEAR or year_to != ANY_YEAR:
|
| 173 |
+
lo = int(year_from) if year_from != ANY_YEAR else YEARS[0]
|
| 174 |
+
hi = int(year_to) if year_to != ANY_YEAR else YEARS[-1]
|
| 175 |
+
year_range = (min(lo, hi), max(lo, hi))
|
| 176 |
+
chunk_range = None
|
| 177 |
+
if chunk_min is not None or chunk_max is not None:
|
| 178 |
+
chunk_range = (int(chunk_min or 0), int(chunk_max if chunk_max is not None else 100000))
|
| 179 |
+
|
| 180 |
+
return {
|
| 181 |
+
"scoring": scoring,
|
| 182 |
+
"bm25": bm25,
|
| 183 |
+
"lucene_mode": lucene_mode,
|
| 184 |
+
"clauses": clauses,
|
| 185 |
+
"raw": raw,
|
| 186 |
+
"text_filters": text_filters,
|
| 187 |
+
"text_filter_join": join,
|
| 188 |
+
"tickers": tickers or [],
|
| 189 |
+
"exclude_tickers": excl or [],
|
| 190 |
+
"year_range": year_range,
|
| 191 |
+
"chunk_range": chunk_range,
|
| 192 |
+
}
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
def builder_preview(top_k, *values):
|
| 196 |
+
try:
|
| 197 |
+
req = fq.build_request(builder_state(list(values)), int(top_k), INCLUDE_FIELDS)
|
| 198 |
+
except fq.QueryError as e:
|
| 199 |
+
return "", "", f"โ ๏ธ {e}"
|
| 200 |
+
return json.dumps(req, indent=2), fq.to_python(req), ""
|
| 201 |
+
|
| 202 |
+
|
| 203 |
+
def builder_values_from_state(state: dict) -> list:
|
| 204 |
+
values = [state["scoring"], state["bm25"], state["lucene_mode"], state["raw"]]
|
| 205 |
+
clauses = [c for c in state["clauses"] if str(c.get("value") or "").strip()]
|
| 206 |
+
for i in range(N_CLAUSES):
|
| 207 |
+
c = clauses[i] if i < len(clauses) else {}
|
| 208 |
+
values += [c.get("occur", "SHOULD"), c.get("kind", "term(s)"), c.get("value", ""),
|
| 209 |
+
c.get("slop", 0), c.get("boost", 1.0)]
|
| 210 |
+
for i in range(N_TEXT_FILTERS):
|
| 211 |
+
f = state["text_filters"][i] if i < len(state["text_filters"]) else {}
|
| 212 |
+
values += [f.get("op", "$match_phrase"), f.get("value", ""), f.get("negate", False)]
|
| 213 |
+
yr = state["year_range"]
|
| 214 |
+
ch = state["chunk_range"]
|
| 215 |
+
values += [
|
| 216 |
+
state["text_filter_join"],
|
| 217 |
+
list(state["tickers"]),
|
| 218 |
+
list(state["exclude_tickers"]),
|
| 219 |
+
str(yr[0]) if yr else ANY_YEAR,
|
| 220 |
+
str(yr[1]) if yr else ANY_YEAR,
|
| 221 |
+
ch[0] if ch else None,
|
| 222 |
+
ch[1] if ch else None,
|
| 223 |
+
]
|
| 224 |
+
return values
|
| 225 |
+
|
| 226 |
+
|
| 227 |
+
def visibility(scoring: str, lucene_mode: str):
|
| 228 |
+
lucene = scoring == "query_string"
|
| 229 |
+
return (
|
| 230 |
+
gr.update(visible=not lucene),
|
| 231 |
+
gr.update(visible=lucene),
|
| 232 |
+
gr.update(visible=lucene and lucene_mode == "clauses"),
|
| 233 |
+
gr.update(visible=lucene and lucene_mode == "raw"),
|
| 234 |
+
)
|
| 235 |
+
|
| 236 |
+
|
| 237 |
+
def load_builder_example(name: str):
|
| 238 |
+
ex = BUILDER_EXAMPLES.get(name)
|
| 239 |
+
state = ex["state"] if ex else fq.BLANK
|
| 240 |
+
info = f"โน๏ธ **Demonstrates:** {ex['shows']}\n\n{ex['description']}" if ex else ""
|
| 241 |
+
return builder_values_from_state(state) + list(visibility(state["scoring"], state["lucene_mode"])) + [info]
|
| 242 |
+
|
| 243 |
+
|
| 244 |
+
def on_lucene_mode(lucene_mode, scoring, raw, *clause_values):
|
| 245 |
+
vis = visibility(scoring, lucene_mode)
|
| 246 |
+
if lucene_mode == "raw" and not str(raw or "").strip():
|
| 247 |
+
rows = [clause_values[i:i + 5] for i in range(0, len(clause_values), 5)]
|
| 248 |
+
clauses = [{"occur": o, "kind": k, "value": v, "slop": s, "boost": b} for o, k, v, s, b in rows]
|
| 249 |
+
try:
|
| 250 |
+
raw = fq.clauses_to_lucene(clauses)
|
| 251 |
+
except fq.QueryError:
|
| 252 |
+
pass
|
| 253 |
+
return (*vis, raw)
|
| 254 |
+
|
| 255 |
+
|
| 256 |
+
def do_builder(top_k, *values):
|
| 257 |
+
state = builder_state(list(values))
|
| 258 |
+
try:
|
| 259 |
+
req = fq.build_request(state, int(top_k), INCLUDE_FIELDS)
|
| 260 |
+
resp = search(req)
|
| 261 |
+
except fq.QueryError as e:
|
| 262 |
+
return error_html(str(e))
|
| 263 |
+
except Exception as e:
|
| 264 |
+
return error_html(f"Search failed: {e}")
|
| 265 |
+
return results_html([match_record(m) for m in resp.matches], fq.highlight_terms(state))
|
| 266 |
+
|
| 267 |
+
|
| 268 |
+
def compare_example(name: str):
|
| 269 |
+
ex = next((e for e in fq.MODE_EXAMPLES["compare"] if e["name"] == name), None)
|
| 270 |
+
n_builder = len(load_builder_example(BLANK))
|
| 271 |
+
if not ex:
|
| 272 |
+
return [gr.update(), gr.update(), ""] + [gr.update()] * n_builder
|
| 273 |
+
builder = load_builder_example(ex["builder_example"]) if ex.get("builder_example") else [gr.update()] * n_builder
|
| 274 |
+
return [ex["values"]["c_query"], ex["values"]["c_source"], f"โน๏ธ {ex['description']}"] + builder
|
| 275 |
+
|
| 276 |
+
|
| 277 |
+
def compare_source_preview(source, top_k, *values):
|
| 278 |
+
if source != "builder":
|
| 279 |
+
return gr.update(visible=False), gr.update(visible=False)
|
| 280 |
+
try:
|
| 281 |
+
req = fq.build_request(builder_state(list(values)), int(top_k), INCLUDE_FIELDS)
|
| 282 |
+
except fq.QueryError as e:
|
| 283 |
+
return gr.update(visible=True, value=f"// Query builder: {e}"), gr.update(visible=False)
|
| 284 |
+
shown = {k: req[k] for k in ("score_by", "filter") if k in req}
|
| 285 |
+
return gr.update(visible=True, value=json.dumps(shown, indent=2)), gr.update(visible="filter" in req)
|
| 286 |
+
|
| 287 |
+
|
| 288 |
+
def do_compare(query, source, share_filters, tickers, years, top_k, *values):
|
| 289 |
+
empty = ("", "", pd.DataFrame(), "", "")
|
| 290 |
+
if not str(query or "").strip():
|
| 291 |
+
return (error_html("Enter a query."),) + empty[1:]
|
| 292 |
+
dense_extra = None
|
| 293 |
+
if source == "builder":
|
| 294 |
+
fts_state = builder_state(list(values))
|
| 295 |
+
if share_filters:
|
| 296 |
+
dense_extra = fq.build_filter(fts_state)
|
| 297 |
+
else:
|
| 298 |
+
fts_state = {**fq.BLANK, "scoring": "text", "bm25": query}
|
| 299 |
+
try:
|
| 300 |
+
res = run_compare(query, fts_state, metadata_filter(tickers, years), dense_extra, int(top_k))
|
| 301 |
+
except fq.QueryError as e:
|
| 302 |
+
return (error_html(f"Query builder: {e}"),) + empty[1:]
|
| 303 |
+
except Exception as e:
|
| 304 |
+
return (error_html(f"Search failed: {e}"),) + empty[1:]
|
| 305 |
+
|
| 306 |
+
stats = compare_stats(res)
|
| 307 |
+
dense, fts, terms = res["dense"], res["fts"], res["terms"]
|
| 308 |
+
dense_rank, fts_rank, shared, union = stats["dense_rank"], stats["fts_rank"], stats["shared"], stats["union"]
|
| 309 |
+
|
| 310 |
+
metrics = "".join([
|
| 311 |
+
metric_html("Overlap", f"{len(shared)} / {stats['k']}", f"Jaccard {len(shared) / len(union):.0%}" if union else ""),
|
| 312 |
+
metric_html("Only dense", str(len(dense_rank.keys() - shared))),
|
| 313 |
+
metric_html("Only full-text", str(len(fts_rank.keys() - shared))),
|
| 314 |
+
metric_html("Dense latency", f"{res['embed_ms'] + res['dense_ms']:.0f} ms",
|
| 315 |
+
f"embed {res['embed_ms']:.0f} + search {res['dense_ms']:.0f}"),
|
| 316 |
+
metric_html("Full-text latency", f"{res['fts_ms']:.0f} ms"),
|
| 317 |
+
])
|
| 318 |
+
summary = f"<div class='metrics'>{metrics}</div>"
|
| 319 |
+
if stats["cov_dense"] is not None:
|
| 320 |
+
summary += (
|
| 321 |
+
f"<div class='count'><b>Keyword coverage</b> โ average share of the full-text terms "
|
| 322 |
+
f"({html.escape(', '.join(terms))}) present in each result: dense <b>{stats['cov_dense']:.0%}</b> vs "
|
| 323 |
+
f"full-text <b>{stats['cov_fts']:.0%}</b>. Low dense coverage means it is finding paraphrases and related "
|
| 324 |
+
"concepts the keywords miss.</div>"
|
| 325 |
+
)
|
| 326 |
+
|
| 327 |
+
def side(rows, other_ranks, other_name, title):
|
| 328 |
+
cards = []
|
| 329 |
+
for i, r in enumerate(rows, 1):
|
| 330 |
+
hits, n = keyword_coverage(r["text"], terms)
|
| 331 |
+
cov = f" ยท keywords {hits}/{n}" if n else ""
|
| 332 |
+
cards.append(card_html(r, terms, rank=i, badge=rank_badge(r["id"], other_ranks, other_name) + cov))
|
| 333 |
+
body = "".join(cards) or "<div class='count'>No results.</div>"
|
| 334 |
+
return f"<div><h4>{title}</h4>{body}</div>"
|
| 335 |
+
|
| 336 |
+
side_by_side = (
|
| 337 |
+
"<div class='cols'>"
|
| 338 |
+
+ side(dense, fts_rank, "full-text", "Dense (semantic)")
|
| 339 |
+
+ side(fts, dense_rank, "dense", "Full-text (BM25 / Lucene)")
|
| 340 |
+
+ "</div>"
|
| 341 |
+
)
|
| 342 |
+
|
| 343 |
+
fused_cards = []
|
| 344 |
+
for i, (doc_id, score) in enumerate(rrf_fuse(stats), 1):
|
| 345 |
+
r = {**stats["by_id"][doc_id], "score": score}
|
| 346 |
+
d, f = dense_rank.get(doc_id), fts_rank.get(doc_id)
|
| 347 |
+
fused_cards.append(card_html(r, terms, rank=i, badge=f"dense #{d if d else 'โ'} ยท full-text #{f if f else 'โ'}"))
|
| 348 |
+
fused = (
|
| 349 |
+
f"<div class='count'>Reciprocal rank fusion merges both lists client-side: score = ฮฃ 1 / ({RRF_K} + rank). "
|
| 350 |
+
"Documents found by both searches rise to the top โ a preview of what a two-query hybrid returns.</div>"
|
| 351 |
+
+ "".join(fused_cards)
|
| 352 |
+
)
|
| 353 |
+
|
| 354 |
+
return summary, side_by_side, pd.DataFrame(rank_rows(res, stats)), fused, json.dumps(res["fts_req"], indent=2)
|
| 355 |
+
|
| 356 |
+
|
| 357 |
+
with gr.Blocks(title="SEC FTS Hybrid Comparison") as app:
|
| 358 |
+
gr.Markdown(
|
| 359 |
+
"# SEC 10-K Search โ Pinecone full-text search vs dense vectors\n"
|
| 360 |
+
"Six companies (AAPL, AMZN, F, GM, MSFT, ORCL) ยท 10-K filings 2019โ2024 ยท ~24k chunks. "
|
| 361 |
+
"Full-text search uses Pinecone's document-schema BM25 / Lucene; dense search uses OpenAI "
|
| 362 |
+
"`text-embedding-3-small`."
|
| 363 |
+
)
|
| 364 |
+
if missing_env():
|
| 365 |
+
gr.Markdown(
|
| 366 |
+
f"โ ๏ธ **Missing environment variable(s): {', '.join(missing_env())}.** "
|
| 367 |
+
"Set them in `.env` locally, or as Secrets in the Space settings."
|
| 368 |
+
)
|
| 369 |
+
|
| 370 |
+
with gr.Sidebar():
|
| 371 |
+
gr.Markdown("### Filters")
|
| 372 |
+
side_tickers = gr.CheckboxGroup(TICKERS, label="Ticker", info="Empty = all companies")
|
| 373 |
+
side_years = gr.CheckboxGroup([str(y) for y in YEARS], label="Year", info="Empty = all years")
|
| 374 |
+
top_k = gr.Slider(3, 25, value=10, step=1, label="Results")
|
| 375 |
+
gr.Markdown("_The Query builder tab has its own filters; these apply to the other tabs._")
|
| 376 |
+
|
| 377 |
+
with gr.Tabs():
|
| 378 |
+
with gr.Tab("Full-text"):
|
| 379 |
+
ft_ex = gr.Dropdown([BLANK] + [e["name"] for e in fq.MODE_EXAMPLES["fulltext"]], value=BLANK, label="Example queries")
|
| 380 |
+
ft_info = gr.Markdown()
|
| 381 |
+
ft_query = gr.Textbox(label="Keyword query", placeholder="e.g. revenue growth operating income")
|
| 382 |
+
ft_btn = gr.Button("Search", variant="primary")
|
| 383 |
+
ft_out = gr.HTML()
|
| 384 |
+
|
| 385 |
+
with gr.Tab("Query builder (full-text)"):
|
| 386 |
+
b_ex = gr.Dropdown([BLANK] + list(BUILDER_EXAMPLES), value=BLANK, label="Example queries",
|
| 387 |
+
info="Load a pre-built complex query into the builder, then tweak it.")
|
| 388 |
+
b_info = gr.Markdown()
|
| 389 |
+
with gr.Accordion("Syntax cheat sheet", open=False):
|
| 390 |
+
gr.Markdown(CHEAT_SHEET)
|
| 391 |
+
|
| 392 |
+
gr.Markdown("### 1 ยท Scoring (ranks results)")
|
| 393 |
+
b_scoring = gr.Radio([("Lucene (query_string)", "query_string"), ("BM25 keywords (text)", "text")],
|
| 394 |
+
value="query_string", show_label=False)
|
| 395 |
+
b_bm25 = gr.Textbox(label="BM25 keywords", placeholder="e.g. battery cells charging range", visible=False)
|
| 396 |
+
b_lucene_mode = gr.Radio([("Clause builder", "clauses"), ("Raw Lucene", "raw")], value="clauses", label="Lucene input")
|
| 397 |
+
with gr.Group() as b_clause_group:
|
| 398 |
+
gr.Markdown("Each row is one clause. **MUST** โ `+`, **MUST NOT** โ `-`, **SHOULD** โ optional (adds score). "
|
| 399 |
+
"Slop applies to phrases; boost multiplies a clause's weight. Empty rows are ignored.")
|
| 400 |
+
clause_comps = []
|
| 401 |
+
for i in range(N_CLAUSES):
|
| 402 |
+
with gr.Row(equal_height=True):
|
| 403 |
+
clause_comps += [
|
| 404 |
+
gr.Dropdown(fq.OCCURS, value="SHOULD", label="Occur", scale=1, min_width=110),
|
| 405 |
+
gr.Dropdown(fq.KINDS, value="term(s)", label="Kind", scale=1, min_width=120),
|
| 406 |
+
gr.Textbox(label="Value", scale=3, min_width=160),
|
| 407 |
+
gr.Number(value=0, label="Slop ~N", precision=0, minimum=0, maximum=50, scale=1, min_width=80),
|
| 408 |
+
gr.Number(value=1.0, label="Boost ^N", minimum=0.1, maximum=20, step=0.5, scale=1, min_width=80),
|
| 409 |
+
]
|
| 410 |
+
b_raw = gr.Textbox(label="Lucene query", lines=3, visible=False,
|
| 411 |
+
placeholder='text:(("supply chain" OR semiconductor) AND shortage) NOT text:(covid)')
|
| 412 |
+
|
| 413 |
+
gr.Markdown("### 2 ยท Filters (hard constraints)")
|
| 414 |
+
gr.Markdown("**Text-match filters** on the `text` field. Empty rows are ignored.")
|
| 415 |
+
tf_comps = []
|
| 416 |
+
for i in range(N_TEXT_FILTERS):
|
| 417 |
+
with gr.Row(equal_height=True):
|
| 418 |
+
tf_comps += [
|
| 419 |
+
gr.Dropdown(fq.MATCH_OPS, value="$match_phrase", label="Operator", scale=1, min_width=140),
|
| 420 |
+
gr.Textbox(label="Value", scale=3, min_width=160),
|
| 421 |
+
gr.Checkbox(label="NOT", value=False, scale=0, min_width=70),
|
| 422 |
+
]
|
| 423 |
+
b_join = gr.Radio([("AND ($and) โ every condition", "all"), ("OR ($or) โ at least one", "any")],
|
| 424 |
+
value="all", label="Combine text-match filters with")
|
| 425 |
+
gr.Markdown("**Metadata filters**")
|
| 426 |
+
with gr.Row():
|
| 427 |
+
b_tickers = gr.CheckboxGroup(TICKERS, label="Ticker is one of ($in)")
|
| 428 |
+
b_excl = gr.CheckboxGroup(TICKERS, label="Ticker is not ($nin)")
|
| 429 |
+
with gr.Row():
|
| 430 |
+
year_choices = [ANY_YEAR] + [str(y) for y in YEARS]
|
| 431 |
+
b_year_from = gr.Dropdown(year_choices, value=ANY_YEAR, label="Year from ($gte)")
|
| 432 |
+
b_year_to = gr.Dropdown(year_choices, value=ANY_YEAR, label="Year to ($lte)")
|
| 433 |
+
b_chunk_min = gr.Number(value=None, label="Chunk index โฅ", precision=0, minimum=0)
|
| 434 |
+
b_chunk_max = gr.Number(value=None, label="Chunk index โค", precision=0, minimum=0)
|
| 435 |
+
|
| 436 |
+
gr.Markdown("### 3 ยท Request")
|
| 437 |
+
b_warn = gr.Markdown()
|
| 438 |
+
with gr.Tabs():
|
| 439 |
+
with gr.Tab("JSON"):
|
| 440 |
+
b_json = gr.Code(language="json", interactive=False, show_label=False)
|
| 441 |
+
with gr.Tab("Python"):
|
| 442 |
+
b_py = gr.Code(language="python", interactive=False, show_label=False)
|
| 443 |
+
b_btn = gr.Button("Search", variant="primary")
|
| 444 |
+
b_out = gr.HTML()
|
| 445 |
+
|
| 446 |
+
with gr.Tab("Semantic (dense)"):
|
| 447 |
+
sem_ex = gr.Dropdown([BLANK] + [e["name"] for e in fq.MODE_EXAMPLES["semantic"]], value=BLANK, label="Example queries")
|
| 448 |
+
sem_info = gr.Markdown()
|
| 449 |
+
sem_query = gr.Textbox(label="Semantic query", placeholder="e.g. risks related to supply chain disruption")
|
| 450 |
+
sem_btn = gr.Button("Search", variant="primary")
|
| 451 |
+
sem_out = gr.HTML()
|
| 452 |
+
|
| 453 |
+
with gr.Tab("Hybrid (semantic + text filter)"):
|
| 454 |
+
hy_ex = gr.Dropdown([BLANK] + [e["name"] for e in fq.MODE_EXAMPLES["hybrid"]], value=BLANK, label="Example queries")
|
| 455 |
+
hy_info = gr.Markdown()
|
| 456 |
+
with gr.Row():
|
| 457 |
+
hy_sem = gr.Textbox(label="Semantic query (drives ranking)", placeholder="e.g. cloud infrastructure investment")
|
| 458 |
+
hy_txt = gr.Textbox(label="Must-contain keywords (full-text filter)", placeholder="e.g. Azure AWS capital expenditure",
|
| 459 |
+
info="All tokens must appear in the chunk ($match_all).")
|
| 460 |
+
hy_btn = gr.Button("Search", variant="primary")
|
| 461 |
+
hy_out = gr.HTML()
|
| 462 |
+
|
| 463 |
+
with gr.Tab("Compare: dense vs full-text"):
|
| 464 |
+
c_ex = gr.Dropdown([BLANK] + [e["name"] for e in fq.MODE_EXAMPLES["compare"]], value=BLANK, label="Example queries")
|
| 465 |
+
c_info = gr.Markdown()
|
| 466 |
+
c_query = gr.Textbox(label="Query", placeholder="e.g. risks from rising interest rates on consumer demand",
|
| 467 |
+
info="Embedded for the dense search. Also used as BM25 keywords unless the full-text side uses the query builder.")
|
| 468 |
+
c_source = gr.Radio([("Same query as BM25 keywords", "same"), ("Current query-builder query", "builder")],
|
| 469 |
+
value="same", label="Full-text side")
|
| 470 |
+
c_builder_preview = gr.Code(language="json", interactive=False, label="Query-builder request (full-text side)", visible=False)
|
| 471 |
+
c_share = gr.Checkbox(label="Apply the builder's filters to the dense search too", value=False, visible=False,
|
| 472 |
+
info="Makes the dense side a hybrid, isolating the effect of the ranking signal.")
|
| 473 |
+
gr.Markdown("_Sidebar ticker / year filters apply to both searches._")
|
| 474 |
+
c_btn = gr.Button("Compare", variant="primary")
|
| 475 |
+
c_summary = gr.HTML()
|
| 476 |
+
with gr.Tabs():
|
| 477 |
+
with gr.Tab("Side by side"):
|
| 478 |
+
c_side = gr.HTML()
|
| 479 |
+
with gr.Tab("Rank comparison"):
|
| 480 |
+
c_table = gr.Dataframe(interactive=False, wrap=True)
|
| 481 |
+
gr.Markdown("_ฮ rank < 0 โ dense ranks it higher; > 0 โ full-text ranks it higher. Blank rank = not in that top-k._")
|
| 482 |
+
with gr.Tab("Fused (RRF)"):
|
| 483 |
+
c_rrf = gr.HTML()
|
| 484 |
+
with gr.Accordion("Full-text request sent", open=False):
|
| 485 |
+
c_req = gr.Code(language="json", interactive=False, show_label=False)
|
| 486 |
+
|
| 487 |
+
side_inputs = [side_tickers, side_years, top_k]
|
| 488 |
+
builder_inputs = [b_scoring, b_bm25, b_lucene_mode, b_raw, *clause_comps, *tf_comps,
|
| 489 |
+
b_join, b_tickers, b_excl, b_year_from, b_year_to, b_chunk_min, b_chunk_max]
|
| 490 |
+
builder_vis = [b_bm25, b_lucene_mode, b_clause_group, b_raw]
|
| 491 |
+
|
| 492 |
+
ft_ex.change(lambda n: mode_example("fulltext", n, ["ft_query"]), ft_ex, [ft_query, ft_info])
|
| 493 |
+
sem_ex.change(lambda n: mode_example("semantic", n, ["sem_query"]), sem_ex, [sem_query, sem_info])
|
| 494 |
+
hy_ex.change(lambda n: mode_example("hybrid", n, ["hy_sem", "hy_txt"]), hy_ex, [hy_sem, hy_txt, hy_info])
|
| 495 |
+
|
| 496 |
+
gr.on([ft_btn.click, ft_query.submit], do_fulltext, [ft_query, *side_inputs], ft_out)
|
| 497 |
+
gr.on([sem_btn.click, sem_query.submit], do_semantic, [sem_query, *side_inputs], sem_out)
|
| 498 |
+
gr.on([hy_btn.click, hy_sem.submit, hy_txt.submit], do_hybrid, [hy_sem, hy_txt, *side_inputs], hy_out)
|
| 499 |
+
|
| 500 |
+
b_ex.change(load_builder_example, b_ex, builder_inputs + builder_vis + [b_info])
|
| 501 |
+
b_scoring.change(visibility, [b_scoring, b_lucene_mode], builder_vis)
|
| 502 |
+
b_lucene_mode.change(on_lucene_mode, [b_lucene_mode, b_scoring, b_raw, *clause_comps], builder_vis + [b_raw])
|
| 503 |
+
gr.on([c.change for c in builder_inputs] + [top_k.change, app.load], builder_preview,
|
| 504 |
+
[top_k, *builder_inputs], [b_json, b_py, b_warn])
|
| 505 |
+
b_btn.click(do_builder, [top_k, *builder_inputs], b_out)
|
| 506 |
+
|
| 507 |
+
c_ex.change(compare_example, c_ex, [c_query, c_source, c_info] + builder_inputs + builder_vis + [b_info])
|
| 508 |
+
gr.on([c_source.change] + [c.change for c in builder_inputs], compare_source_preview,
|
| 509 |
+
[c_source, top_k, *builder_inputs], [c_builder_preview, c_share])
|
| 510 |
+
gr.on([c_btn.click, c_query.submit], do_compare, [c_query, c_source, c_share, *side_inputs, *builder_inputs],
|
| 511 |
+
[c_summary, c_side, c_table, c_rrf, c_req])
|
| 512 |
+
|
| 513 |
+
demo = app
|
| 514 |
+
|
| 515 |
+
if __name__ == "__main__":
|
| 516 |
+
app.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False, css=CSS, theme=gr.themes.Soft())
|
fts_queries.py
ADDED
|
@@ -0,0 +1,386 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import json
|
| 2 |
+
import re
|
| 3 |
+
|
| 4 |
+
TEXT_FIELD = "text"
|
| 5 |
+
|
| 6 |
+
OCCURS = ["SHOULD", "MUST", "MUST NOT"]
|
| 7 |
+
OCCUR_PREFIX = {"SHOULD": "", "MUST": "+", "MUST NOT": "-"}
|
| 8 |
+
KINDS = ["term(s)", "phrase", "phrase prefix"]
|
| 9 |
+
MATCH_OPS = ["$match_phrase", "$match_all", "$match_any"]
|
| 10 |
+
|
| 11 |
+
LUCENE_KEYWORDS = {"and", "or", "not", "text"}
|
| 12 |
+
STOP_WORDS = {
|
| 13 |
+
"a", "an", "the", "of", "on", "in", "to", "for", "from", "by", "with", "at", "as", "is", "are",
|
| 14 |
+
"was", "be", "it", "its", "this", "that", "these", "those", "our", "we", "their", "about", "into",
|
| 15 |
+
}
|
| 16 |
+
|
| 17 |
+
|
| 18 |
+
class QueryError(ValueError):
|
| 19 |
+
pass
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
def clause_to_lucene(c: dict) -> str | None:
|
| 23 |
+
value = str(c.get("value") or "").strip().replace('"', "")
|
| 24 |
+
if not value:
|
| 25 |
+
return None
|
| 26 |
+
kind = c.get("kind") or "term(s)"
|
| 27 |
+
slop = int(c.get("slop") or 0)
|
| 28 |
+
boost = float(c.get("boost") or 1.0)
|
| 29 |
+
|
| 30 |
+
if kind == "phrase":
|
| 31 |
+
expr = f'"{value}"' + (f"~{slop}" if slop > 0 else "")
|
| 32 |
+
elif kind == "phrase prefix":
|
| 33 |
+
if len(value.split()) < 2:
|
| 34 |
+
raise QueryError(
|
| 35 |
+
f"Phrase prefix '{value}' needs at least two words "
|
| 36 |
+
'(e.g. "artificial intel") โ single-term wildcards are not supported.'
|
| 37 |
+
)
|
| 38 |
+
expr = f'"{value}"*'
|
| 39 |
+
else:
|
| 40 |
+
expr = value
|
| 41 |
+
|
| 42 |
+
prefix = OCCUR_PREFIX.get(c.get("occur") or "SHOULD", "")
|
| 43 |
+
if kind == "term(s)" and len(value.split()) > 1 and (prefix or boost != 1.0):
|
| 44 |
+
expr = f"({value})"
|
| 45 |
+
if boost != 1.0:
|
| 46 |
+
expr += f"^{boost:g}"
|
| 47 |
+
return prefix + expr
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def clauses_to_lucene(clauses: list[dict]) -> str:
|
| 51 |
+
parts, positive = [], 0
|
| 52 |
+
for c in clauses:
|
| 53 |
+
expr = clause_to_lucene(c)
|
| 54 |
+
if expr is None:
|
| 55 |
+
continue
|
| 56 |
+
parts.append(expr)
|
| 57 |
+
if c.get("occur") != "MUST NOT":
|
| 58 |
+
positive += 1
|
| 59 |
+
if not parts:
|
| 60 |
+
return ""
|
| 61 |
+
if positive == 0:
|
| 62 |
+
raise QueryError("Add at least one SHOULD or MUST clause โ a query of only exclusions is rejected.")
|
| 63 |
+
return f"{TEXT_FIELD}:({' '.join(parts)})"
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def build_filter(state: dict) -> dict | None:
|
| 67 |
+
conds = []
|
| 68 |
+
|
| 69 |
+
text_conds = []
|
| 70 |
+
for f in state.get("text_filters") or []:
|
| 71 |
+
value = str(f.get("value") or "").strip()
|
| 72 |
+
op = f.get("op")
|
| 73 |
+
if not value or op not in MATCH_OPS:
|
| 74 |
+
continue
|
| 75 |
+
cond = {TEXT_FIELD: {op: value}}
|
| 76 |
+
text_conds.append({"$not": cond} if f.get("negate") else cond)
|
| 77 |
+
if len(text_conds) > 1 and state.get("text_filter_join") == "any":
|
| 78 |
+
conds.append({"$or": text_conds})
|
| 79 |
+
else:
|
| 80 |
+
conds.extend(text_conds)
|
| 81 |
+
|
| 82 |
+
if state.get("tickers"):
|
| 83 |
+
conds.append({"ticker": {"$in": list(state["tickers"])}})
|
| 84 |
+
if state.get("exclude_tickers"):
|
| 85 |
+
conds.append({"ticker": {"$nin": list(state["exclude_tickers"])}})
|
| 86 |
+
if state.get("year_range"):
|
| 87 |
+
lo, hi = state["year_range"]
|
| 88 |
+
conds.append({"year": {"$gte": lo, "$lte": hi}})
|
| 89 |
+
if state.get("chunk_range"):
|
| 90 |
+
lo, hi = state["chunk_range"]
|
| 91 |
+
conds.append({"chunk_index": {"$gte": lo, "$lte": hi}})
|
| 92 |
+
|
| 93 |
+
if not conds:
|
| 94 |
+
return None
|
| 95 |
+
return conds[0] if len(conds) == 1 else {"$and": conds}
|
| 96 |
+
|
| 97 |
+
|
| 98 |
+
def build_score_by(state: dict) -> list[dict]:
|
| 99 |
+
if state["scoring"] == "text":
|
| 100 |
+
query = (state.get("bm25") or "").strip()
|
| 101 |
+
if not query:
|
| 102 |
+
raise QueryError("Enter BM25 keywords.")
|
| 103 |
+
return [{"type": "text", "field": TEXT_FIELD, "query": query}]
|
| 104 |
+
|
| 105 |
+
if state.get("lucene_mode") == "raw":
|
| 106 |
+
query = (state.get("raw") or "").strip()
|
| 107 |
+
else:
|
| 108 |
+
query = clauses_to_lucene(state.get("clauses") or [])
|
| 109 |
+
if not query:
|
| 110 |
+
raise QueryError("The Lucene query is empty.")
|
| 111 |
+
return [{"type": "query_string", "query": query}]
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
def build_request(state: dict, top_k: int, include_fields: list[str]) -> dict:
|
| 115 |
+
req = {
|
| 116 |
+
"namespace": "__default__",
|
| 117 |
+
"top_k": top_k,
|
| 118 |
+
"score_by": build_score_by(state),
|
| 119 |
+
"include_fields": include_fields,
|
| 120 |
+
}
|
| 121 |
+
filt = build_filter(state)
|
| 122 |
+
if filt:
|
| 123 |
+
req["filter"] = filt
|
| 124 |
+
return req
|
| 125 |
+
|
| 126 |
+
|
| 127 |
+
def to_python(req: dict) -> str:
|
| 128 |
+
lines = ["resp = idx.documents.search("]
|
| 129 |
+
for key, value in req.items():
|
| 130 |
+
rendered = json.dumps(value, indent=4).replace("\n", "\n ")
|
| 131 |
+
lines.append(f" {key}={rendered},")
|
| 132 |
+
lines.append(")")
|
| 133 |
+
return "\n".join(lines)
|
| 134 |
+
|
| 135 |
+
|
| 136 |
+
def highlight_terms(state: dict) -> list[str]:
|
| 137 |
+
sources = []
|
| 138 |
+
if state["scoring"] == "text":
|
| 139 |
+
sources.append(state.get("bm25") or "")
|
| 140 |
+
elif state.get("lucene_mode") == "raw":
|
| 141 |
+
raw = state.get("raw") or ""
|
| 142 |
+
raw = re.sub(r"-\"[^\"]*\"|-\w+|NOT\s+\w+:\([^)]*\)|NOT\s+\"[^\"]*\"|NOT\s+\w+", " ", raw)
|
| 143 |
+
sources.append(raw)
|
| 144 |
+
else:
|
| 145 |
+
sources += [str(c.get("value") or "") for c in state.get("clauses") or [] if c.get("occur") != "MUST NOT"]
|
| 146 |
+
sources += [str(f.get("value") or "") for f in state.get("text_filters") or [] if not f.get("negate")]
|
| 147 |
+
|
| 148 |
+
terms = set()
|
| 149 |
+
for s in sources:
|
| 150 |
+
for word in re.findall(r"[A-Za-z][A-Za-z0-9']+", s):
|
| 151 |
+
if word.lower() not in LUCENE_KEYWORDS | STOP_WORDS:
|
| 152 |
+
terms.add(word.lower())
|
| 153 |
+
return sorted(terms, key=len, reverse=True)
|
| 154 |
+
|
| 155 |
+
|
| 156 |
+
def _clause(value, occur="SHOULD", kind="term(s)", slop=0, boost=1.0):
|
| 157 |
+
return {"occur": occur, "kind": kind, "value": value, "slop": slop, "boost": boost}
|
| 158 |
+
|
| 159 |
+
|
| 160 |
+
def _tf(op, value, negate=False):
|
| 161 |
+
return {"op": op, "value": value, "negate": negate}
|
| 162 |
+
|
| 163 |
+
|
| 164 |
+
BLANK = {
|
| 165 |
+
"scoring": "query_string",
|
| 166 |
+
"bm25": "",
|
| 167 |
+
"lucene_mode": "clauses",
|
| 168 |
+
"clauses": [_clause("")],
|
| 169 |
+
"raw": "",
|
| 170 |
+
"text_filters": [],
|
| 171 |
+
"text_filter_join": "all",
|
| 172 |
+
"tickers": [],
|
| 173 |
+
"exclude_tickers": [],
|
| 174 |
+
"year_range": None,
|
| 175 |
+
"chunk_range": None,
|
| 176 |
+
}
|
| 177 |
+
|
| 178 |
+
|
| 179 |
+
def _example(**overrides) -> dict:
|
| 180 |
+
return {**BLANK, **overrides}
|
| 181 |
+
|
| 182 |
+
|
| 183 |
+
EXAMPLES = [
|
| 184 |
+
{
|
| 185 |
+
"name": "Supply-chain shortages, excluding COVID",
|
| 186 |
+
"shows": "Boolean nesting โ AND / OR / NOT with a phrase inside a group",
|
| 187 |
+
"description": "Chunks that talk about a shortage of either the supply chain or semiconductors, "
|
| 188 |
+
"but never mention COVID. Parentheses group the OR before the AND.",
|
| 189 |
+
"state": _example(
|
| 190 |
+
lucene_mode="raw",
|
| 191 |
+
raw='text:(("supply chain" OR semiconductor) AND shortage) NOT text:(covid)',
|
| 192 |
+
),
|
| 193 |
+
},
|
| 194 |
+
{
|
| 195 |
+
"name": "Cyber incidents (required / excluded / boosted)",
|
| 196 |
+
"shows": "Clause builder โ MUST (+), MUST NOT (โ), SHOULD with a boost",
|
| 197 |
+
"description": "Every result must mention cybersecurity; ransomware mentions count triple, "
|
| 198 |
+
"breach is a nice-to-have, and chunks about insurance are dropped.",
|
| 199 |
+
"state": _example(
|
| 200 |
+
clauses=[
|
| 201 |
+
_clause("cybersecurity", occur="MUST"),
|
| 202 |
+
_clause("ransomware", boost=3.0),
|
| 203 |
+
_clause("breach"),
|
| 204 |
+
_clause("insurance", occur="MUST NOT"),
|
| 205 |
+
],
|
| 206 |
+
),
|
| 207 |
+
},
|
| 208 |
+
{
|
| 209 |
+
"name": "Rising interest rates (proximity)",
|
| 210 |
+
"shows": 'Phrase slop โ "interest rates increase"~5',
|
| 211 |
+
"description": "The three words must appear within 5 positions of each other, so "
|
| 212 |
+
'"increases in interest rates" and "interest rates may continue to increase" both match. '
|
| 213 |
+
"Stemming lets increase match increases/increased.",
|
| 214 |
+
"state": _example(
|
| 215 |
+
clauses=[_clause("interest rates increase", kind="phrase", slop=5)],
|
| 216 |
+
),
|
| 217 |
+
},
|
| 218 |
+
{
|
| 219 |
+
"name": "China trade & tariffs (term boost)",
|
| 220 |
+
"shows": "Term boosting โ tariffs^3 outweighs trade and china",
|
| 221 |
+
"description": "BM25 over three terms, but tariffs is weighted 3ร. Compare with the boost "
|
| 222 |
+
"removed to see ranking shift toward generic trade/China text.",
|
| 223 |
+
"state": _example(
|
| 224 |
+
clauses=[
|
| 225 |
+
_clause("tariffs", boost=3.0),
|
| 226 |
+
_clause("trade"),
|
| 227 |
+
_clause("china"),
|
| 228 |
+
],
|
| 229 |
+
),
|
| 230 |
+
},
|
| 231 |
+
{
|
| 232 |
+
"name": "AI mentions (phrase prefix)",
|
| 233 |
+
"shows": 'Phrase prefix โ "artificial intel"* matches intelligence / intelligent',
|
| 234 |
+
"description": "The last word is matched as a prefix. Combined with a required cloud term "
|
| 235 |
+
"and restricted to 2022+ filings. Note: prefix matches all share a constant score.",
|
| 236 |
+
"state": _example(
|
| 237 |
+
clauses=[
|
| 238 |
+
_clause("artificial intel", occur="MUST", kind="phrase prefix"),
|
| 239 |
+
_clause("cloud", occur="MUST"),
|
| 240 |
+
],
|
| 241 |
+
year_range=(2022, 2024),
|
| 242 |
+
),
|
| 243 |
+
},
|
| 244 |
+
{
|
| 245 |
+
"name": "EV batteries at Ford & GM, 2022+",
|
| 246 |
+
"shows": "BM25 ranking + $match_phrase hard filter + metadata $in / range",
|
| 247 |
+
"description": 'Plain BM25 over battery terms, restricted to chunks containing the exact phrase '
|
| 248 |
+
'"electric vehicles", from Ford or GM, filed 2022โ2024.',
|
| 249 |
+
"state": _example(
|
| 250 |
+
scoring="text",
|
| 251 |
+
bm25="battery cells charging range",
|
| 252 |
+
text_filters=[_tf("$match_phrase", "electric vehicles")],
|
| 253 |
+
tickers=["f", "gm"],
|
| 254 |
+
year_range=(2022, 2024),
|
| 255 |
+
),
|
| 256 |
+
},
|
| 257 |
+
{
|
| 258 |
+
"name": "Cloud growth, no pandemic talk, tech only",
|
| 259 |
+
"shows": "$not + $match_any exclusion filter, ticker $nin",
|
| 260 |
+
"description": "Ranks on cloud revenue growth but excludes any chunk mentioning COVID or "
|
| 261 |
+
"pandemic, and excludes the automakers entirely.",
|
| 262 |
+
"state": _example(
|
| 263 |
+
clauses=[_clause("cloud revenue growth")],
|
| 264 |
+
text_filters=[_tf("$match_any", "covid pandemic", negate=True)],
|
| 265 |
+
exclude_tickers=["f", "gm"],
|
| 266 |
+
),
|
| 267 |
+
},
|
| 268 |
+
{
|
| 269 |
+
"name": "Regulators: EC or DOJ",
|
| 270 |
+
"shows": "$or across two $match_phrase filters",
|
| 271 |
+
"description": "Antitrust / litigation language, limited to chunks that name either the "
|
| 272 |
+
"European Commission or the Department of Justice.",
|
| 273 |
+
"state": _example(
|
| 274 |
+
clauses=[_clause("antitrust", boost=2.0), _clause("litigation"), _clause("investigation")],
|
| 275 |
+
text_filters=[
|
| 276 |
+
_tf("$match_phrase", "european commission"),
|
| 277 |
+
_tf("$match_phrase", "department of justice"),
|
| 278 |
+
],
|
| 279 |
+
text_filter_join="any",
|
| 280 |
+
),
|
| 281 |
+
},
|
| 282 |
+
{
|
| 283 |
+
"name": "Inflation & input costs, 2021โ2023",
|
| 284 |
+
"shows": "$match_all (all tokens, any order) + year range",
|
| 285 |
+
"description": "Ranks on inflation; every result must contain raw, materials and costs somewhere "
|
| 286 |
+
"in the chunk (any order), from the 2021โ2023 filings.",
|
| 287 |
+
"state": _example(
|
| 288 |
+
scoring="text",
|
| 289 |
+
bm25="inflation inflationary pressures",
|
| 290 |
+
text_filters=[_tf("$match_all", "raw materials costs")],
|
| 291 |
+
year_range=(2021, 2023),
|
| 292 |
+
),
|
| 293 |
+
},
|
| 294 |
+
{
|
| 295 |
+
"name": "Buybacks vs. dividends (required groups)",
|
| 296 |
+
"shows": "Required OR-groups โ +(a OR b) +(c OR d)",
|
| 297 |
+
"description": "Each result must contain a buyback term AND a dividend term; the phrase "
|
| 298 |
+
'"share repurchase" is boosted so it outranks bare "repurchase".',
|
| 299 |
+
"state": _example(
|
| 300 |
+
lucene_mode="raw",
|
| 301 |
+
raw='text:(+("share repurchase"^2 OR buyback OR repurchase) +(dividend OR dividends))',
|
| 302 |
+
),
|
| 303 |
+
},
|
| 304 |
+
]
|
| 305 |
+
|
| 306 |
+
|
| 307 |
+
MODE_EXAMPLES = {
|
| 308 |
+
"fulltext": [
|
| 309 |
+
{
|
| 310 |
+
"name": "Buybacks & dividends",
|
| 311 |
+
"description": "BM25 token-OR: chunks with more (and rarer) of these terms rank higher.",
|
| 312 |
+
"values": {"ft_query": "share repurchase program dividends"},
|
| 313 |
+
},
|
| 314 |
+
{
|
| 315 |
+
"name": "Chip shortage hits production",
|
| 316 |
+
"description": "Surfaces the automakers' 2021โ2022 semiconductor supply disclosures.",
|
| 317 |
+
"values": {"ft_query": "semiconductor shortage production"},
|
| 318 |
+
},
|
| 319 |
+
{
|
| 320 |
+
"name": "FX risk",
|
| 321 |
+
"description": "Market-risk language about foreign currency exposure.",
|
| 322 |
+
"values": {"ft_query": "foreign currency exchange rate risk"},
|
| 323 |
+
},
|
| 324 |
+
],
|
| 325 |
+
"semantic": [
|
| 326 |
+
{
|
| 327 |
+
"name": "Interest-rate exposure (question)",
|
| 328 |
+
"description": "A natural-language question โ dense search matches meaning, not exact words.",
|
| 329 |
+
"values": {"sem_query": "How is the company exposed to rising interest rates?"},
|
| 330 |
+
},
|
| 331 |
+
{
|
| 332 |
+
"name": "Supplier concentration in Asia",
|
| 333 |
+
"description": "Filings rarely use this wording; embeddings find single-source and "
|
| 334 |
+
"outsourcing-partner risk disclosures anyway.",
|
| 335 |
+
"values": {"sem_query": "Concerns about dependence on a small number of suppliers in Asia"},
|
| 336 |
+
},
|
| 337 |
+
{
|
| 338 |
+
"name": "AI & data-center buildout",
|
| 339 |
+
"description": "Conceptual query that spans capex, infrastructure and AI discussion.",
|
| 340 |
+
"values": {"sem_query": "Investments in artificial intelligence and data center capacity"},
|
| 341 |
+
},
|
| 342 |
+
],
|
| 343 |
+
"hybrid": [
|
| 344 |
+
{
|
| 345 |
+
"name": "Cloud growth that names Azure",
|
| 346 |
+
"description": "Semantic ranking on cloud growth, but only chunks containing \"Azure\".",
|
| 347 |
+
"values": {"hy_sem": "growth of cloud infrastructure services", "hy_txt": "Azure"},
|
| 348 |
+
},
|
| 349 |
+
{
|
| 350 |
+
"name": "EV transition that names Ultium",
|
| 351 |
+
"description": "GM's battery platform as a hard keyword; ranking by the EV-transition concept.",
|
| 352 |
+
"values": {"hy_sem": "transition to electric vehicles and battery supply", "hy_txt": "Ultium"},
|
| 353 |
+
},
|
| 354 |
+
{
|
| 355 |
+
"name": "Competition probes by the EC",
|
| 356 |
+
"description": "Regulatory-investigation meaning, restricted to chunks mentioning European Commission.",
|
| 357 |
+
"values": {
|
| 358 |
+
"hy_sem": "regulatory investigations into competition practices",
|
| 359 |
+
"hy_txt": "European Commission",
|
| 360 |
+
},
|
| 361 |
+
},
|
| 362 |
+
],
|
| 363 |
+
"compare": [
|
| 364 |
+
{
|
| 365 |
+
"name": "Interest rates vs consumer demand",
|
| 366 |
+
"description": "Same words on both sides: BM25 rewards literal term hits, dense rewards the idea.",
|
| 367 |
+
"values": {"c_query": "risks from rising interest rates on consumer demand", "c_source": "same"},
|
| 368 |
+
},
|
| 369 |
+
{
|
| 370 |
+
"name": "Paraphrase: overseas manufacturing partners",
|
| 371 |
+
"description": "Wording filings don't use verbatim โ expect low overlap and low keyword "
|
| 372 |
+
"coverage on the dense side, where it still finds outsourcing-risk passages.",
|
| 373 |
+
"values": {
|
| 374 |
+
"c_query": "the company depends on a few manufacturing partners overseas",
|
| 375 |
+
"c_source": "same",
|
| 376 |
+
},
|
| 377 |
+
},
|
| 378 |
+
{
|
| 379 |
+
"name": "Lucene builder query vs dense",
|
| 380 |
+
"description": "Loads the \"Supply-chain shortages, excluding COVID\" builder query as the "
|
| 381 |
+
"full-text side against a dense search for the same topic.",
|
| 382 |
+
"values": {"c_query": "semiconductor chip shortage impact on production", "c_source": "builder"},
|
| 383 |
+
"builder_example": "Supply-chain shortages, excluding COVID",
|
| 384 |
+
},
|
| 385 |
+
],
|
| 386 |
+
}
|
requirements.txt
ADDED
|
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
pinecone==10.0.0
|
| 2 |
+
openai==3.22.1
|
| 3 |
+
pandas==3.0.6
|
| 4 |
+
python-dotenv==1.2.3
|
| 5 |
+
spaces==0.51.3
|
search_core.py
ADDED
|
@@ -0,0 +1,237 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
import html
|
| 2 |
+
import os
|
| 3 |
+
import re
|
| 4 |
+
import time
|
| 5 |
+
from concurrent.futures import ThreadPoolExecutor
|
| 6 |
+
from functools import lru_cache
|
| 7 |
+
|
| 8 |
+
from dotenv import load_dotenv
|
| 9 |
+
from openai import OpenAI
|
| 10 |
+
from pinecone import Pinecone
|
| 11 |
+
|
| 12 |
+
import fts_queries as fq
|
| 13 |
+
|
| 14 |
+
load_dotenv()
|
| 15 |
+
|
| 16 |
+
INDEX_NAME = "sec-fts"
|
| 17 |
+
NAMESPACE = "__default__"
|
| 18 |
+
EMBED_MODEL = "text-embedding-3-small"
|
| 19 |
+
|
| 20 |
+
TICKERS = ["aapl", "amzn", "f", "gm", "msft", "orcl"]
|
| 21 |
+
YEARS = list(range(2019, 2025))
|
| 22 |
+
INCLUDE_FIELDS = ["text", "ticker", "filing_type", "year", "chunk_index"]
|
| 23 |
+
RRF_K = 60
|
| 24 |
+
|
| 25 |
+
|
| 26 |
+
def missing_env() -> list[str]:
|
| 27 |
+
return [k for k in ("PINECONE_API_KEY", "OPENAI_API_KEY") if not os.environ.get(k)]
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
@lru_cache(maxsize=1)
|
| 31 |
+
def get_index():
|
| 32 |
+
return Pinecone(api_key=os.environ["PINECONE_API_KEY"]).index(name=INDEX_NAME)
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
@lru_cache(maxsize=1)
|
| 36 |
+
def get_openai() -> OpenAI:
|
| 37 |
+
return OpenAI(api_key=os.environ["OPENAI_API_KEY"])
|
| 38 |
+
|
| 39 |
+
|
| 40 |
+
@lru_cache(maxsize=512)
|
| 41 |
+
def _embed(text: str) -> tuple[float, ...]:
|
| 42 |
+
resp = get_openai().embeddings.create(model=EMBED_MODEL, input=[text])
|
| 43 |
+
return tuple(resp.data[0].embedding)
|
| 44 |
+
|
| 45 |
+
|
| 46 |
+
def embed(text: str) -> list[float]:
|
| 47 |
+
return list(_embed(text))
|
| 48 |
+
|
| 49 |
+
|
| 50 |
+
def metadata_filter(tickers: list[str], years: list[int]) -> dict | None:
|
| 51 |
+
filt = {}
|
| 52 |
+
if tickers:
|
| 53 |
+
filt["ticker"] = {"$in": list(tickers)}
|
| 54 |
+
if years:
|
| 55 |
+
filt["year"] = {"$in": [int(y) for y in years]}
|
| 56 |
+
return filt or None
|
| 57 |
+
|
| 58 |
+
|
| 59 |
+
def and_filters(*filters: dict | None) -> dict | None:
|
| 60 |
+
parts = [f for f in filters if f]
|
| 61 |
+
if not parts:
|
| 62 |
+
return None
|
| 63 |
+
return parts[0] if len(parts) == 1 else {"$and": parts}
|
| 64 |
+
|
| 65 |
+
|
| 66 |
+
def search(req: dict):
|
| 67 |
+
return get_index().documents.search(**req)
|
| 68 |
+
|
| 69 |
+
|
| 70 |
+
def dense_request(query: str, top_k: int, filt: dict | None = None) -> dict:
|
| 71 |
+
req = {
|
| 72 |
+
"namespace": NAMESPACE,
|
| 73 |
+
"top_k": top_k,
|
| 74 |
+
"score_by": [{"type": "dense_vector", "field": "embedding", "values": embed(query)}],
|
| 75 |
+
"include_fields": INCLUDE_FIELDS,
|
| 76 |
+
}
|
| 77 |
+
if filt:
|
| 78 |
+
req["filter"] = filt
|
| 79 |
+
return req
|
| 80 |
+
|
| 81 |
+
|
| 82 |
+
def run_search_text(query: str, tickers: list, years: list, top_k: int):
|
| 83 |
+
req = {
|
| 84 |
+
"namespace": NAMESPACE,
|
| 85 |
+
"top_k": top_k,
|
| 86 |
+
"score_by": [{"type": "text", "field": "text", "query": query}],
|
| 87 |
+
"include_fields": INCLUDE_FIELDS,
|
| 88 |
+
}
|
| 89 |
+
filt = metadata_filter(tickers, years)
|
| 90 |
+
if filt:
|
| 91 |
+
req["filter"] = filt
|
| 92 |
+
return search(req)
|
| 93 |
+
|
| 94 |
+
|
| 95 |
+
def run_search_semantic(query: str, tickers: list, years: list, top_k: int):
|
| 96 |
+
return search(dense_request(query, top_k, metadata_filter(tickers, years)))
|
| 97 |
+
|
| 98 |
+
|
| 99 |
+
def run_search_hybrid(semantic_query: str, text_filter: str, tickers: list, years: list, top_k: int):
|
| 100 |
+
filt = metadata_filter(tickers, years) or {}
|
| 101 |
+
if text_filter.strip():
|
| 102 |
+
filt["text"] = {"$match_all": text_filter.strip()}
|
| 103 |
+
return search(dense_request(semantic_query, top_k, filt or None))
|
| 104 |
+
|
| 105 |
+
|
| 106 |
+
def term_pattern(terms: list[str]) -> re.Pattern | None:
|
| 107 |
+
if not terms:
|
| 108 |
+
return None
|
| 109 |
+
stems = [re.escape(t[:-2] if len(t) > 5 else t) for t in terms]
|
| 110 |
+
return re.compile(r"\b(" + "|".join(stems) + r")\w*", re.IGNORECASE)
|
| 111 |
+
|
| 112 |
+
|
| 113 |
+
def highlight(text: str, terms: list[str]) -> tuple[str, str]:
|
| 114 |
+
escaped = html.escape(text)
|
| 115 |
+
pattern = term_pattern(terms)
|
| 116 |
+
if not pattern:
|
| 117 |
+
return escaped[:300] + ("..." if len(escaped) > 300 else ""), escaped
|
| 118 |
+
marked = pattern.sub(lambda m: f"<mark>{m.group(0)}</mark>", escaped)
|
| 119 |
+
|
| 120 |
+
first = pattern.search(escaped)
|
| 121 |
+
start = max(0, first.start() - 150) if first else 0
|
| 122 |
+
snippet_raw = escaped[start : start + 450]
|
| 123 |
+
snippet = pattern.sub(lambda m: f"<mark>{m.group(0)}</mark>", snippet_raw)
|
| 124 |
+
snippet = ("..." if start else "") + snippet + ("..." if start + 450 < len(escaped) else "")
|
| 125 |
+
return snippet, marked
|
| 126 |
+
|
| 127 |
+
|
| 128 |
+
def keyword_coverage(text: str, terms: list[str]) -> tuple[int, int]:
|
| 129 |
+
hits = sum(1 for t in terms if term_pattern([t]).search(text))
|
| 130 |
+
return hits, len(terms)
|
| 131 |
+
|
| 132 |
+
|
| 133 |
+
def match_record(m) -> dict:
|
| 134 |
+
d = m.to_dict()
|
| 135 |
+
return {
|
| 136 |
+
"id": getattr(m, "_id", d.get("_id")),
|
| 137 |
+
"score": getattr(m, "_score", getattr(m, "score", d.get("_score", 0.0))),
|
| 138 |
+
"ticker": str(d.get("ticker", "?")).upper(),
|
| 139 |
+
"year": int(d.get("year", 0)),
|
| 140 |
+
"filing": str(d.get("filing_type", "")).upper(),
|
| 141 |
+
"chunk": int(d["chunk_index"]) if d.get("chunk_index") is not None else "?",
|
| 142 |
+
"text": str(d.get("text", "")),
|
| 143 |
+
}
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
def _timed(fn):
|
| 147 |
+
start = time.perf_counter()
|
| 148 |
+
result = fn()
|
| 149 |
+
return result, (time.perf_counter() - start) * 1000
|
| 150 |
+
|
| 151 |
+
|
| 152 |
+
def run_compare(query: str, fts_state: dict, base_filter: dict | None, dense_extra_filter: dict | None, top_k: int) -> dict:
|
| 153 |
+
fts_req = fq.build_request(fts_state, top_k, INCLUDE_FIELDS)
|
| 154 |
+
fts_req["filter"] = and_filters(fts_req.get("filter"), base_filter)
|
| 155 |
+
if not fts_req["filter"]:
|
| 156 |
+
del fts_req["filter"]
|
| 157 |
+
|
| 158 |
+
def dense():
|
| 159 |
+
_, embed_ms = _timed(lambda: embed(query))
|
| 160 |
+
resp, search_ms = _timed(lambda: search(dense_request(query, top_k, and_filters(base_filter, dense_extra_filter))))
|
| 161 |
+
return resp, embed_ms, search_ms
|
| 162 |
+
|
| 163 |
+
with ThreadPoolExecutor(max_workers=2) as pool:
|
| 164 |
+
dense_future = pool.submit(dense)
|
| 165 |
+
fts_future = pool.submit(_timed, lambda: search(fts_req))
|
| 166 |
+
dense_resp, embed_ms, dense_ms = dense_future.result()
|
| 167 |
+
fts_resp, fts_ms = fts_future.result()
|
| 168 |
+
|
| 169 |
+
return {
|
| 170 |
+
"dense": [match_record(m) for m in dense_resp.matches],
|
| 171 |
+
"fts": [match_record(m) for m in fts_resp.matches],
|
| 172 |
+
"embed_ms": embed_ms,
|
| 173 |
+
"dense_ms": dense_ms,
|
| 174 |
+
"fts_ms": fts_ms,
|
| 175 |
+
"fts_req": fts_req,
|
| 176 |
+
"terms": fq.highlight_terms(fts_state),
|
| 177 |
+
}
|
| 178 |
+
|
| 179 |
+
|
| 180 |
+
def compare_stats(res: dict) -> dict:
|
| 181 |
+
dense, fts, terms = res["dense"], res["fts"], res["terms"]
|
| 182 |
+
dense_rank = {r["id"]: i for i, r in enumerate(dense, 1)}
|
| 183 |
+
fts_rank = {r["id"]: i for i, r in enumerate(fts, 1)}
|
| 184 |
+
shared = dense_rank.keys() & fts_rank.keys()
|
| 185 |
+
union = dense_rank.keys() | fts_rank.keys()
|
| 186 |
+
|
| 187 |
+
def avg_coverage(rows):
|
| 188 |
+
if not rows or not terms:
|
| 189 |
+
return None
|
| 190 |
+
return sum(keyword_coverage(r["text"], terms)[0] / len(terms) for r in rows) / len(rows)
|
| 191 |
+
|
| 192 |
+
return {
|
| 193 |
+
"dense_rank": dense_rank,
|
| 194 |
+
"fts_rank": fts_rank,
|
| 195 |
+
"shared": shared,
|
| 196 |
+
"union": union,
|
| 197 |
+
"by_id": {r["id"]: r for r in dense + fts},
|
| 198 |
+
"cov_dense": avg_coverage(dense),
|
| 199 |
+
"cov_fts": avg_coverage(fts),
|
| 200 |
+
"k": max(len(dense), len(fts)),
|
| 201 |
+
}
|
| 202 |
+
|
| 203 |
+
|
| 204 |
+
def rank_rows(res: dict, stats: dict) -> list[dict]:
|
| 205 |
+
terms = res["terms"]
|
| 206 |
+
rows = []
|
| 207 |
+
for doc_id in stats["union"]:
|
| 208 |
+
r = stats["by_id"][doc_id]
|
| 209 |
+
d, f = stats["dense_rank"].get(doc_id), stats["fts_rank"].get(doc_id)
|
| 210 |
+
hits, n = keyword_coverage(r["text"], terms)
|
| 211 |
+
rows.append({
|
| 212 |
+
"ticker": r["ticker"],
|
| 213 |
+
"year": r["year"],
|
| 214 |
+
"chunk": r["chunk"],
|
| 215 |
+
"dense rank": d,
|
| 216 |
+
"full-text rank": f,
|
| 217 |
+
"ฮ rank (dense โ FTS)": d - f if d and f else None,
|
| 218 |
+
"found by": "both" if d and f else ("dense" if d else "full-text"),
|
| 219 |
+
"keywords": f"{hits}/{n}" if n else "",
|
| 220 |
+
"snippet": r["text"][:140].replace("\n", " "),
|
| 221 |
+
})
|
| 222 |
+
rows.sort(key=lambda x: min(x["dense rank"] or 999, x["full-text rank"] or 999))
|
| 223 |
+
return rows
|
| 224 |
+
|
| 225 |
+
|
| 226 |
+
def rrf_fuse(stats: dict) -> list[tuple[str, float]]:
|
| 227 |
+
fused = {}
|
| 228 |
+
for ranks in (stats["dense_rank"], stats["fts_rank"]):
|
| 229 |
+
for doc_id, rank in ranks.items():
|
| 230 |
+
fused[doc_id] = fused.get(doc_id, 0.0) + 1 / (RRF_K + rank)
|
| 231 |
+
return sorted(fused.items(), key=lambda kv: kv[1], reverse=True)[: stats["k"]]
|
| 232 |
+
|
| 233 |
+
|
| 234 |
+
def rank_badge(doc_id: str, other_ranks: dict, other_name: str) -> str:
|
| 235 |
+
if doc_id in other_ranks:
|
| 236 |
+
return f"๐ also #{other_ranks[doc_id]} in {other_name}"
|
| 237 |
+
return f"โ only in {'dense' if other_name == 'full-text' else 'full-text'}"
|