Search Alchemy

import bm25s
import pandas as pd

# Step 1: Create BM25 index (one-time setup)
df = pd.read_parquet("data/amazon_products.parquet")
corpus = df["FullText"].tolist()
corpus_tokens = bm25s.tokenize(corpus, stopwords="en")

retriever = bm25s.BM25()
retriever.index(corpus_tokens)
retriever.save("data/bm25_index")

# Step 2: Load index and search
bm25_index = bm25s.BM25.load("data/bm25_index", load_corpus=False)
query_tokens = bm25s.tokenize(query, stopwords="en")
results, scores = bm25_index.retrieve(query_tokens, k=5)

# Extract top results
top_products = [df.iloc[idx] for idx in results[0]]

from openai import OpenAI
import faiss
import numpy as np

client = OpenAI(
    base_url="https://api.fireworks.ai/inference/v1"
)

# Generate embeddings
response = client.embeddings.create(
    model="accounts/fireworks/models/qwen3-embedding-8b",
    input=[query] + documents
)

# Extract embeddings
query_emb = np.array(response.data[0].embedding)
doc_embs = np.array([d.embedding for d in response.data[1:]])

# FAISS search
index = faiss.IndexFlatIP(doc_embs.shape[1])
index.add(doc_embs)
scores, indices = index.search(query_emb.reshape(1, -1), k=5)

# Query expansion with LLM
response = client.chat.completions.create(
    model="accounts/fireworks/models/llama-v3p1-8b-instruct",
    messages=[{
        "role": "user",
        "content": f"Extract 2-3 key search concepts from: {query}"
    }]
)

expanded_query = response.choices[0].message.content

# Search with expanded query
response = client.embeddings.create(
    model="accounts/fireworks/models/qwen3-embedding-8b",
    input=[expanded_query] + documents
)

# Continue with embedding search...

# First get top 20 candidates from Stage 3
top_20_results = get_stage_3_results(query, k=20)

# Rerank with Fireworks reranker
rerank_response = client.post(
    "https://api.fireworks.ai/inference/v1/rerank",
    json={
        "model": "fireworks/qwen3-reranker-8b",
        "query": query,
        "documents": [r["text"] for r in top_20_results],
        "top_n": 5
    }
)

# Get final ranked results
final_results = [
    top_20_results[r["index"]]
    for r in rerank_response.json()["results"]
]