Metadata-Version: 2.5
Name: annotations-vector-search
Version: 0.1.0
Summary: Dataset generation from RavenPack Annotations using Snowflake vector search
Author-email: Alberto Martinez <amartinezg@ravenpack.com>
License-Expression: LicenseRef-Proprietary
License-File: LICENSE
Requires-Python: >=3.10
Requires-Dist: pyarrow>=23.0.1
Requires-Dist: pydantic-settings>=2.15.0
Requires-Dist: pydantic>=2.14.0
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: sentence-transformers>=6.1.0
Requires-Dist: snowflake-connector-python[pandas]>=4.8.0
Requires-Dist: torch>=2.14.1
Requires-Dist: typer>=0.27.3
Description-Content-Type: text/markdown

# annotations-vector-search

Generate datasets from RavenPack Annotations using Snowflake vector search.

Give it a time window, a list of companies (RP_ENTITY_IDs) and one or more search texts.
It finds the text blocks of news, transcripts and other documents that mention those
companies and are most similar in meaning to your texts, and writes them as a dataset
(Parquet or CSV): one row per block, company and search text, with its similarity, text,
headline and the company's mentions. It needs access to RavenPack Annotations in Snowflake
and a local copy of the RavenPack embedding model (see "Setup" and "Getting the model").

## Quick start

```bash
pip install annotations-vector-search   # see "Setup" for the CPU-only torch build
```

Set the Snowflake credentials and the model path, as environment variables or in a `.env`
file in the working directory:

```bash
SF_ACCOUNT=...
SF_USER=...
SF_PWD=...
SF_WAREHOUSE=...
RP_EMBEDDINGS_MODEL_PATH=/path/to/model
```

Describe the search in `params.yaml`:

```yaml
providers: [DJ]
start_date: 2026-09-01
end_date: 2026-09-05
entities: [D8442A]
search_texts:
  - Company gained share in its core market at the expense of competitors
percentile: 90              # keep the top 10% most similar blocks
```

Run it from the command line:

```bash
avs run params.yaml --out data/run_001   # data/run_001/data/part-*.parquet + manifest.json
```

or from Python:

```python
from annotations_vector_search import SearchParams, search

df = search(SearchParams.from_file("params.yaml"))  # a pandas DataFrame
```

## How it works

```
SearchParams ──► Embedder (local model) ──► SQL (filters + VECTOR_COSINE_SIMILARITY) ──► Snowflake ──► Parquet/CSV + manifest
```

Data model (`RAVENPACK_ANNOTATIONS_2_0.RP_ANNOTATIONS_<provider>`, e.g. `RP_ANNOTATIONS_DJ`):

| Table | Grain | Used for |
|---|---|---|
| `DOCUMENT_ENTITY` | document × entity | entity filter, name, relevance, sentiment |
| `ENTITY_ANNOTATION` | block × entity | blocks where the entity is mentioned |
| `DOCUMENT_BLOCK` | text block | `BLOCK_TEXT`, `BLOCK_INDEX` |
| `DOCUMENT_EMBEDDINGS_1_0` or `DOCUMENT_EMBEDDINGS_1_0_NEW` | text block | `EMBEDDING VECTOR(FLOAT, 384)` |

The first three tables have the same name for every provider; the embeddings table depends on
the provider (`PROVIDER_EMBEDDING_TABLES` in `config.py`).

Several providers are searched in one query: each provider's tables are read in its own
`UNION ALL` branch (so each keeps its partition pruning), and rows are only ever joined within
the same provider. Repeated texts are collapsed across providers too: the earliest occurrence
is kept, and on equal `TIMESTAMP_UTC` the first provider alphabetically.

A block matches an entity when the entity is mentioned in that block (`ENTITY_ANNOTATION`).
All tables are clustered on the month of `TIMESTAMP_UTC`, so the query filters each one on the
time window and joins on `(RP_DOCUMENT_ID, TIMESTAMP_UTC, BLOCK_ID)` to keep partition pruning.
Repeated texts (e.g. the same story syndicated by several sources) are collapsed before
scoring: for the same entity and `BLOCK_TEXT`, only the earliest occurrence is kept (across
providers too, see above). Similarity is then computed once per remaining block.

Transcript providers (`FACTSET_TRANSCRIPTS`, `QUARTR`, `SP_GLOBAL_TRANSCRIPTS`; `CHAIN_KEYS` in
`query_builder.py`) are collapsed by document chain instead of by text: for the same entity and
`DOCUMENT_ANALYTICS.PROVIDER_DOCUMENT_CHAIN_ID`, only the blocks of the earliest document are
kept (on equal `TIMESTAMP_UTC`, the smallest `RP_DOCUMENT_ID`). For `QUARTR` the chain is the
event, the first `:`-separated part of the id, so its filings, audio, slides and transcript
are one chain. These providers take no part in the text deduplication. The chain is returned
as `document_chain_id`, and with `chunk_days` the chunks' results are collapsed again by
chain once downloaded, as a single query would. Transcript documents are also kept only when
they come on the day of their event or the day after: `DATEDIFF('day',
TRANSCRIPT_METADATA.TRANSCRIPT_EVENT_DATETIME, TIMESTAMP_UTC) BETWEEN 0 AND 1`
(`DOCUMENT_FILTERS` in `query_builder.py`).

## Setup

The package runs the embedding model on the CPU only, so install the CPU build of torch: on
Linux the default build from PyPI also bundles CUDA (several GB). Python packages can't name
the index their dependencies come from, so this is up to the installing command.

With pip, install torch from the PyTorch CPU index first:

```bash
pip install torch --index-url https://download.pytorch.org/whl/cpu
pip install annotations-vector-search
```

With uv:

```bash
uv pip install annotations-vector-search --torch-backend=cpu
```

In a uv project (`uv add annotations-vector-search`), point torch to the CPU index in your
`pyproject.toml`:

```toml
[tool.uv.sources]
torch = { index = "pytorch-cpu", marker = "sys_platform == 'linux'" }

[[tool.uv.index]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
explicit = true
```

On macOS and Windows the default torch build is already CPU-only.

Settings come from environment variables, or from a `.env` file in the working directory
(where you run `avs` or your script). Create one with the Snowflake credentials and the
model path (see "Getting the model"):

```bash
# --- Snowflake connection ---
SF_ACCOUNT=
SF_USER=
SF_PWD=
SF_ROLE=
SF_WAREHOUSE=
# SF_AUTHENTICATOR=externalbrowser
# SF_PRIVATE_KEY_FILE=
# SF_PRIVATE_KEY_FILE_PWD=

# --- Embedding model (local directory, see README "Getting the model") ---
RP_EMBEDDINGS_MODEL_PATH=

# --- Annotations database (default shown; the provider schema is set per search) ---
# RP_ANNOTATIONS_DATABASE=RAVENPACK_ANNOTATIONS_2_0

# --- Guards ---
# RP_MAX_WINDOW_DAYS=366

# --- Integration tests (Snowflake DEV) ---
# RP_TEST_ENTITY=
# RP_TEST_START=2024-01-01
# RP_TEST_END=2024-01-08
```

See "Configuration" for what each setting does.

## Getting the model

The embedding model is **not** part of this package or repository, and the package never
downloads it. Get a copy of the model directory and point the package to it:

1. Request the model from RavenPack customer support. You receive a model directory
   (~130 MB): keep it as it is, it must be the version that produced the annotation
   embeddings stored in Snowflake.
2. Set `RP_EMBEDDINGS_MODEL_PATH=/path/to/model` in `.env` or your shell, or pass
   `--model-path` / `model_path=`.

Write search texts as plain sentences; the package embeds them the way the stored blocks
were embedded.

## Usage

A search is described by a params file (YAML or JSON). Save this as `params.yaml` and
replace the entities and search texts with your own:

```yaml
providers: [QUARTR, DJ]     # DJ, WEB, SEC_FILINGS, ... (see the table below)
start_date: 2026-09-01      # UTC, inclusive
end_date: 2026-09-05        # UTC, exclusive
entities:                   # RP_ENTITY_IDs, or instead a CSV: entities_file: entities.csv
  - D8442A
search_texts:
  - Company gained share in its core market at the expense of competitors
# min_relevance: 50         # optional: minimum ENTITY_RELEVANCE (1-100)
percentile: 80              # keep blocks at or above this percentile of the candidates' scores (0-100)
percentile_scope: per_text  # per_text | per_entity: scores the percentile is computed over
# threshold: 0.5            # optional: minimum cosine similarity, on top of the percentile
chunk_days: 31              # optional: split the window into sub-queries
warehouse_size: SMALL       # optional: resize the warehouse for the run, then restore it
parallel_queries: 1         # optional: run up to this many chunk queries at the same time
```

The fields:

| Field | Description |
|---|---|
| `start_date`, `end_date` | UTC window, end exclusive. Required, max `RP_MAX_WINDOW_DAYS` (366) |
| `providers` | Annotations providers, searched together (default `[DJ]`; a single `provider: DJ` also works): `AN`, `BZG`, `DJ`, `DJSN`, `FACTSET_TRANSCRIPTS`, `FINANCIAL_TIMES`, `FLY`, `FXS`, `MT`, `PODCHASER`, `PUBT`, `QUARTR`, `SEC_FILINGS`, `SP_GLOBAL_TRANSCRIPTS`, `THIRD_BRIDGE`, `WEB`, `WEB_SECONDARY` |
| `entities` or `entities_file` | RP_ENTITY_IDs, as a list or as a CSV file (see below) |
| `search_texts` | One or more texts, searched in a single query (see the size note below) |
| `min_relevance` | Optional minimum `ENTITY_RELEVANCE` (1-100) |
| `percentile`, `percentile_scope` | Keep the blocks scoring at or above this percentile (0-100) of the requested entities' block scores: per search text (`per_text`, default) or per entity and search text (`per_entity`). E.g. `99` keeps the top 1% (see below) |
| `threshold` | Optional minimum cosine similarity, applied on top of `percentile`. At least one of the two is required |
| `chunk_days` | Split long windows into sub-queries, merged client side (see the note below) |
| `parallel_queries` | Run up to this many chunk queries, and download up to this many results, at the same time (default 1; see below) |
| `warehouse_size` | Optional: run the queries on a warehouse of this size (`X-SMALL` … `6X-LARGE`, see below) |

**How `percentile` selects.** Cosine similarities of this model sit in a narrow band (a good
match is often only ~0.6, unrelated news ~0.2-0.5) and the band differs per search text, so a
fixed `threshold` is hard to choose. `percentile` instead ranks each block against the other
candidates: every block of the requested entities in the window is scored, and those at or
above the percentile of those scores are kept.

- `per_text`: one cutoff per search text, over the blocks of all requested entities (a
  block mentioning several of them counts once).
- `per_entity`: one cutoff per entity and search text, so every entity keeps its own top
  slice, including entities with little coverage.
- It is relative: some blocks are always returned, even when none is really relevant. Add a
  `threshold` as an absolute floor if needed.
- With `chunk_days`, the cutoff is computed within each chunk (an exact percentile over the
  whole window would need every score, i.e. scoring twice).

**Note on `chunk_days`.** A run with the same params always returns the same rows, but the
result depends on where the chunk boundaries fall, so changing `chunk_days` (or `start_date`,
which moves the boundaries) can change it:

- Repeated texts are collapsed within each chunk only. A text that appears in two chunks is
  scored and returned once per chunk, while an unchunked run keeps only the earliest
  occurrence. With several providers this also applies across providers. Document chains
  (transcript providers) are the exception: they are collapsed again across chunks once the
  results are downloaded, so those rows are held in memory until the end and written last.
- The percentile cutoff is computed per chunk, so it changes with the chunk boundaries.
- Rows of several chunks are grouped by `sentence`, in window order: each chunk's rows are
  best first, but the output has no similarity to order them across chunks.

To compare runs, keep `chunk_days` and the window the same, or leave `chunk_days` unset
when the window is small enough for a single query.

`avs run` prints how long the run took, and `manifest.json` keeps the breakdown in
`timings`: `embed_seconds` (model load + embedding), `snowflake_seconds` (connection + all
queries, on the resized warehouse), `query_seconds` (one per chunk), `download_seconds`
(downloading and converting the results, after the warehouse is restored), `write_seconds`
(writing the files) and `total_seconds`. With `parallel_queries` above 1 the `query_seconds`
overlap, so they no longer add up to `snowflake_seconds`. `avs run` writes while it
downloads, so `write_seconds` is part of `download_seconds`.

### Downloading results

A run has two phases. First every chunk query runs (up to `parallel_queries` at once) and
its result stays stored in Snowflake; nothing is downloaded yet. As soon as the last query
finishes, the warehouse gets its size back (see "Warehouse size"). Then the results are
downloaded, which needs no warehouse.

Each stored result is read as Arrow batches, up to `parallel_queries` results at a time,
with 10 download threads each. The queries send a block's `chunk_text` only on its first row
and a document's `headline` only on its first row (the rows of a block are adjacent), and
the client fills them back in: a block's text would otherwise be downloaded once per
matching search text and entity, about 10 times on the benchmark.

### Parallel queries

With `chunk_days`, the chunk queries run one after another by default. `parallel_queries: N`
submits up to N of them at once (asynchronously, on the same session, so the same
`QUERY_TAG` and warehouse) and starts the next as each finishes. Results are merged in window
order, so they are the same as with `parallel_queries: 1`. If a chunk fails, the chunks still
running are cancelled and its error is raised. The results are then downloaded up to N at a
time too (see "Downloading results").

Chunks running at the same time share the warehouse, so the gain comes mostly from
overlapping each query's compilation and queueing, not from more compute: combine it with a
bigger `warehouse_size` for heavy runs. A warehouse runs up to `MAX_CONCURRENCY_LEVEL` queries
at once (8 by default) and queues the rest, so higher values bring nothing. Smaller chunks
also mean more repeated texts across chunks (see the note on `chunk_days`).

### Warehouse size

With `warehouse_size`, the run looks up the size of the session's warehouse (`SF_WAREHOUSE`
or the user's default), resizes it to `warehouse_size` and waits until the new size is
ready, runs the queries, and sets the previous size back as soon as they have finished,
before the results are downloaded (the download needs no warehouse, so it isn't billed for
it), also when a query fails. If
the warehouse already has that size, nothing is changed. Sizes are case-insensitive and
Snowflake's other spellings work too (`xsmall`, `XXLARGE`, `x4large`, ...).

Things to keep in mind:

- `ALTER WAREHOUSE` needs the `MODIFY` privilege on the warehouse.
- The new size applies to everyone using that warehouse while the queries run, not only to
  this run. Two runs resizing the same warehouse at the same time can leave it at the wrong
  size, since each one restores the size it saw when it started.
- If the process is killed (not just a failed query), the size is not restored.
- `--dry-run` and `avs sql` never resize.

### Entities from a CSV

Instead of `entities`, set `entities_file` to a CSV of RP_ENTITY_IDs, e.g. a universe of
companies:

```csv
id
D8442A
228D42
```

Like any other file path, a relative path is resolved against the working directory (where
you run `avs`), not the folder of the params file:

```yaml
entities_file: entities.csv                     # relative to the working directory
entities_file: ~/data/universes/entities.csv    # ~ is expanded
entities_file: ${RP_DATA_DIR}/entities.csv      # so are environment variables
```

Or pass it on the command line with `--entities-file`. It replaces the params file's
`entities` / `entities_file`, so one params file can be run over several company lists:

```bash
avs run params.yaml --entities-file entities.csv --out data/run_002
```

If the file is missing, the error shows the full path that was looked for.

The file needs a header row: the IDs come from the `id` column (case-insensitive, other
columns are ignored). Without it, loading fails and the error shows the header found.
Blank lines and duplicates are dropped. The file is read when the params are loaded, so the
manifest lists the entity IDs actually searched, not the file path.

### CLI

```bash
avs sql params.yaml                      # SQL template + binds, no model
avs estimate params.yaml                 # rows + Parquet size per percentile, before a run
avs run params.yaml --out data/run_001   # data/part-*.parquet + manifest.json
avs run params.yaml --out data/run_001 --max-file-mb 256  # smaller parts
avs run params.yaml --out data/run_001 --dry-run  # embed + print runnable SQL
```

### Estimating a run

`avs estimate` tells how big a run would be before running it, in under a minute and
without the vector search:

```
percentile            rows  parquet size
         0     289,472,160       65.5 GB
        90      28,947,216        5.4 GB
        99       2,894,722      513.2 MB
      99.9         289,472       63.7 MB
```

- **Rows**: it counts the candidates, the requested entities' (block, entity) mentions in the
  window (`DOCUMENT_ENTITY` x `ENTITY_ANNOTATION`, with the same filters as the search), per
  provider. At percentile p a search returns about candidates x search texts x (1 - p/100)
  rows.
- **Size**: it samples `--sample` candidate rows (default 5,000) with their real text,
  headline and detections, simulates which search texts each kept block matches, and writes
  them with the real writer (a block's rows together, as a run returns them).
- Both are **upper bounds**. The search also drops repeated texts and blocks without an
  embedding, and blocks relevant to one search text tend to match several, which compresses
  better than the estimate assumes. On 7 days of the benchmark at `percentile: 50`: 2.04M
  rows and 18.7 MB estimated, 1.78M rows and 13 MB real.
- `--percentiles 90,99,99.9` picks the rows of the table. The params' `percentile` and
  `threshold` are ignored; `warehouse_size` is used.

### Python

```python
from annotations_vector_search import SearchParams, run, run_to_dir, search

params = SearchParams.from_file("params.yaml")
df = search(params)                      # DataFrame
result = run(params)                     # DataFrame + manifest (query ids, model, SQL hash, timings)
manifest = run_to_dir(params, "data/run_001")  # streamed to files, like `avs run`
```

Every CLI command has a Python equivalent. The params can also be built in code, and the
model and the Snowflake connection loaded once and reused across calls (`embedder=`,
`connection=`, `settings=`); otherwise each call loads the model and connects again:

```python
from datetime import datetime

from annotations_vector_search import SearchParams, SearchSettings, run_to_dir, search
from annotations_vector_search.embedder import DIM
from annotations_vector_search.estimate import estimate
from annotations_vector_search.pipeline import build_queries, check_window, load_embedder
from annotations_vector_search.snowflake.connection import connect
from annotations_vector_search.snowflake.query_builder import build_query

params = SearchParams(
    providers=["QUARTR", "DJ"],
    start_date=datetime(2026, 9, 1),  # UTC, inclusive
    end_date=datetime(2026, 9, 5),  # UTC, exclusive
    entities=["D8442A"],  # or entities_file="entities.csv"
    search_texts=["Company gained share in its core market at the expense of competitors"],
    percentile=80,
    chunk_days=31,
    warehouse_size="SMALL",
)
settings = SearchSettings()  # RP_* variables / .env
check_window(params, settings)

# `avs sql`: the SQL with placeholder vectors (no model, no Snowflake)
print(build_query(params, [[0.0] * DIM] * len(params.search_texts), settings).sql)

embedder = load_embedder(None, settings)  # RP_EMBEDDINGS_MODEL_PATH

# `avs run --dry-run`: SQL with the real query vectors, runnable in Snowsight
for query in build_queries(params, embedder, settings):
    print(query.render() + ";")

with connect() as conn:  # SF_* variables / .env
    # `avs estimate`
    est = estimate(params, percentiles=[50, 80, 90, 99], connection=conn, settings=settings)
    for e in est.percentiles:
        print(f"p{e.percentile:g}: {e.rows:,} rows, {e.parquet_bytes / 2**20:.1f} MB")

    # In memory
    df = search(params, embedder=embedder, connection=conn, settings=settings)

    # `avs run`
    manifest = run_to_dir(
        params, "data/run_001", embedder=embedder, connection=conn, settings=settings
    )
    print(f"{manifest.row_count} rows, {manifest.timings.total_seconds:.1f}s")
```

### Output files

```
data/run_001/
  manifest.json           # params, model, query ids, timings, files, row_count
  data/
    part-00000.parquet    # about --max-file-mb each (default 512)
    part-00001.parquet
```

`avs run` (and `run_to_dir`) streams the results to disk while they are downloaded,
so memory use stays at a few row groups however large the result: rows are
buffered 250,000 at a time, written as one Parquet row group, and a new part is started once
the current one reaches `--max-file-mb` (a part can exceed it by up to one row group). Read
all parts as one dataset with `pd.read_parquet("data/run_001/data")` (or Polars, DuckDB,
Spark on the same folder). `--format csv` writes `part-*.csv` files, each with a header.

The output directory must not already contain parts. If a run fails, the parts written so
far stay, but there is no `manifest.json`: treat that directory as incomplete.

`run` keeps the whole result in memory instead; `write_dataset(result.data, result.manifest,
out_dir)` saves it in the same layout.

Output columns (lowercase, the same names as the smart_batching datasets):

| Column | From |
|---|---|
| `sentence` | the search text |
| `similarity` | cosine similarity between the block embedding and the search text |
| `date` | `TIMESTAMP_UTC` (UTC date) |
| `timestamp_utc` | `TIMESTAMP_UTC` |
| `doc_id` | `RP_DOCUMENT_ID` |
| `headline` | `DOCUMENT_ANALYTICS.TITLE` |
| `source_name` | the provider (`DJ`, `BZG`, ...) |
| `document_chain_id` | the document chain (see above) for transcript providers, null for the others |
| `chunk_index` | `BLOCK_INDEX` (position of the block in the document) |
| `chunk_text` | `BLOCK_TEXT` |
| `chunk_sentiment` | `DOCUMENT_BLOCK.SENTIMENT` |
| `entity_id` | `RP_ENTITY_ID` (the entity the row's `detections` belong to) |
| `entity_name` | `DOCUMENT_ENTITY.ENTITY_NAME` (the entity the row's `detections` belong to) |
| `detections` | the entity's mentions in `chunk_text`: a list of `{id, start, end, type}`, `id` the `RP_ENTITY_ID`, `start`/`end` character offsets (end exclusive), `type` `entity`. In Parquet a `list<struct>`, in CSV JSON text |

One row per (block, entity, sentence): a block mentioning several requested entities
appears once per entity, each row with that entity's `detections`. The queries return the
rows ordered by block (provider, document, block index), then sentence and entity.

Parquet parts are ZSTD-compressed. Each row group is sorted by `doc_id`, `chunk_index`,
`sentence`, which puts the repeated `chunk_text` of a block next to each other and compresses
much better (`write_dataset` sorts the whole result). CSV output keeps the query order. With
several results downloaded at once, their rows may be interleaved across the parts.

### How many search texts per query

All query vectors are sent as one JSON bind that the connector inlines into the SQL text,
so the statement grows by about 5 KB per search text (384 values each). Each value is written
as the shortest decimal that reads back as the same float32 (at most 9 significant digits):
the embeddings are stored as float32 (`VECTOR(FLOAT, 384)`), so nothing is lost, while the
default float64 repr would take ~8.5 KB per text. Snowsight displays the stored values with 6
decimals, but they have full float32 precision.

Cost also grows linearly: every candidate block is scored against every search text. Keep it
to a few dozen texts per query; for many more, split them across runs.

## Configuration

All settings come from environment variables or `.env` (see "Setup"):

- `SF_*`: Snowflake connection (`SF_ACCOUNT`, `SF_USER`, `SF_PWD`, `SF_ROLE`, `SF_WAREHOUSE`,
  optional `SF_AUTHENTICATOR`, `SF_PRIVATE_KEY_FILE`).
- `RP_ANNOTATIONS_DATABASE`: database holding the provider schemas (default
  `RAVENPACK_ANNOTATIONS_2_0`). The provider itself is set per search in the params.
- `RP_MAX_WINDOW_DAYS`: cost guard on the window length.

Queries are tagged with `QUERY_TAG = annotations_vector_search:<run_id>` for cost tracking.