Integrating SQL analytics with vector and full-text search

pip install infino pyarrow pandas sentence-transformers datasets

import pyarrow as pa
from datasets import load_dataset
from sentence_transformers import SentenceTransformer
import infino

ds = load_dataset("banking77", split="train")
labels = ds.features["label"].names
texts = ds["text"]
topics = [labels[i].replace("_", " ") for i in ds["label"]]

model = SentenceTransformer("all-MiniLM-L6-v2")   # 384 dimensions
vecs = model.encode(texts, batch_size=256)

db = infino.connect("./bankhelp-sql")   # or "s3://bucket/prefix"

questions = db.create_table(
    "questions",
    pa.schema([
        pa.field("text", pa.large_utf8(), nullable=False),
        pa.field("topic", pa.large_utf8(), nullable=False),
        pa.field("embedding", pa.list_(pa.float32(), 384), nullable=False),
    ]),
    infino.IndexSpec().fts("text").fts("topic").vector("embedding", 384, "cosine"),
)
questions.append([
    {"text": t, "topic": g, "embedding": v.tolist()}
    for t, g, v in zip(texts, topics, vecs)
])


def team_for(topic):
    if any(w in topic for w in ("card", "atm", "pin")):
        return "cards"
    if "transfer" in topic or "beneficiary" in topic:
        return "transfers"
    if "top up" in topic or "topping up" in topic:
        return "top-ups"
    if any(w in topic for w in ("exchange", "currency", "fiat")):
        return "currency"
    return "accounts"


teams = db.create_table(
    "teams",
    pa.schema([
        pa.field("topic", pa.large_utf8(), nullable=False),
        pa.field("team", pa.large_utf8(), nullable=False),
    ]),
    infino.IndexSpec(),
)
teams.append([
    {"topic": l.replace("_", " "), "team": team_for(l.replace("_", " "))}
    for l in labels
])

def vec(text):
    return ",".join(str(x) for x in model.encode(text).tolist())


def show(sql):
    print(db.query_sql(sql).to_pandas().to_string(index=False))

complaint = "my card payment was declined at the shop"
show(f"""
SELECT t.team, COUNT(*) AS questions
FROM vector_search('questions', 'embedding', '{vec(complaint)}', 200) AS v
JOIN teams AS t ON t.topic = v.topic
GROUP BY t.team
ORDER BY questions DESC
""")

     team  questions
    cards        182
 accounts          9
transfers          7
  top-ups          2

show("""
SELECT t.team, COUNT(*) AS mentions
FROM bm25_search('questions', 'text', 'fee', 1000) AS b
JOIN teams AS t ON t.topic = b.topic
GROUP BY t.team
ORDER BY mentions DESC
""")

     team  mentions
    cards       153
 accounts       141
transfers       114
 currency        33

hq = "refund for a transfer that never arrived"
show(f"""
SELECT team, text, score
FROM (
  SELECT t.team, h.text, h.score,
         ROW_NUMBER() OVER (PARTITION BY t.team ORDER BY h.score DESC) AS rank
  FROM hybrid_search('questions', 'text', '{hq}', 'embedding', '{vec(hq)}', 100) AS h
  JOIN teams AS t ON t.topic = h.topic
) AS ranked
WHERE rank = 1
ORDER BY score DESC
""")

     team                                                                    text    score
transfers                                              My transfer hasn't arrived 0.030090
 accounts How can I get a refund for an item I purchased but has not yet arrived? 0.022618
    cards                                                  My card never arrived. 0.016393

show(f"""
SELECT h.text, h.topic, h.score
FROM hybrid_search('questions', 'text', '{hq}', 'embedding', '{vec(hq)}', 100) AS h
JOIN teams AS t ON t.topic = h.topic
WHERE t.team = 'transfers'
ORDER BY h.score DESC
LIMIT 5
""")

                                                 text                                   topic    score
                           My transfer hasn't arrived      transfer not received by recipient 0.030090
                         My transfer has not arrived.      transfer not received by recipient 0.029851
     Money that I have transferred hasn't arrived yet balance not updated after bank transfer 0.025978
The money that I have transferred hasn't arrived yet. balance not updated after bank transfer 0.022918
                 Why has my transfer not arrived yet?      transfer not received by recipient 0.022160

dq = "why was I charged extra"
show(f"""
WITH sides AS (
  SELECT CASE
           WHEN v._id IS NULL THEN 'keyword only'
           WHEN k._id IS NULL THEN 'meaning only'
           ELSE 'both'
         END AS found_by
  FROM bm25_search('questions', 'text', '{dq}', 100) AS k
  FULL OUTER JOIN vector_search('questions', 'embedding', '{vec(dq)}', 100) AS v
    ON v._id = k._id
)
SELECT found_by, COUNT(*) AS questions
FROM sides
GROUP BY found_by
ORDER BY questions DESC
""")

    found_by  questions
meaning only         57
keyword only         57
        both         43