Text rule trends

#!/usr/bin/env python3
"""Embed every item's text with a sentence model (bge-small by default).

Usage:
  python3 embed.py SAMPLE_DIR [--model BAAI/bge-small-en-v1.5]
                   [--backend fastembed|sentence-transformers] [--cache-dir DIR] [--max-chars 300]

Reads SAMPLE_DIR/rules.json, writes SAMPLE_DIR/emb.npy (float32, one row per item, same order).
emb.npy is a local working file; it is not meant to be shared or bundled.
Before embedding: code spans become the word CODE, links keep their text, bold markers go, and
the text is cut at --max-chars characters.

Backends:
  fastembed (default). "BAAI/bge-small-en-v1.5" is served as
      Qdrant's quantised ONNX export. Any model fastembed lists works. pip install fastembed
  sentence-transformers: any Hugging Face sentence model, normalised. pip install sentence-transformers
      Vectors differ slightly from fastembed's for the same model, so clusters will move.
"""
import argparse
import json
import os
import re

import numpy as np


def clean(t, n):
    t = re.sub(r"`[^`]*`", "CODE", t)
    t = re.sub(r"\[([^\]]*)\]\([^)]*\)", r"\1", t)
    t = re.sub(r"\*\*([^*]+)\*\*", r"\1", t)
    return t.strip()[:n]


def main():
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("sample")
    ap.add_argument("--model", default="BAAI/bge-small-en-v1.5")
    ap.add_argument("--backend", default="fastembed", choices=["fastembed", "sentence-transformers"])
    ap.add_argument("--cache-dir", default=None, help="where the model is downloaded")
    ap.add_argument("--batch", type=int, default=256)
    ap.add_argument("--max-chars", type=int, default=300)
    a = ap.parse_args()

    life = json.load(open(os.path.join(a.sample, "rules.json")))["life"]
    texts = [clean(r["text"], a.max_chars) for r in life]
    if a.backend == "fastembed":
        from fastembed import TextEmbedding
        model = TextEmbedding(a.model, cache_dir=a.cache_dir)
        E = np.array(list(model.embed(texts, batch_size=a.batch)), dtype=np.float32)
    else:
        from sentence_transformers import SentenceTransformer
        model = SentenceTransformer(a.model, cache_folder=a.cache_dir)
        E = model.encode(texts, batch_size=a.batch, normalize_embeddings=True, show_progress_bar=False).astype(np.float32)
    np.save(os.path.join(a.sample, "emb.npy"), E)
    print(f"embedded {len(texts)} items with {a.model} -> {E.shape}")


if __name__ == "__main__":
    main()