Гибридный поиск в Pinecone и pgvector

from pinecone import Pinecone, ServerlessSpec

pc = Pinecone(api_key='YOUR_API_KEY')

# Create a sparse-dense index
pc.create_index(
    name='hybrid-index',
    dimension=1536,          # dense vector dimension
    metric='dotproduct',     # must use dotproduct for hybrid
    spec=ServerlessSpec(cloud='aws', region='us-east-1'),
)

index = pc.Index('hybrid-index')

from transformers import AutoTokenizer, AutoModelForMaskedLM
import torch

# Load SPLADE model
tokenizer = AutoTokenizer.from_pretrained('naver/splade-cocondenser-ensembledistil')
model = AutoModelForMaskedLM.from_pretrained('naver/splade-cocondenser-ensembledistil')

def generate_sparse_vector(text: str) -> dict:
    tokens = tokenizer(text, return_tensors='pt', truncation=True, max_length=512)
    with torch.no_grad():
        output = model(**tokens)
    logits = output.logits
    # Max-pool over sequence length and apply log1p activation
    sparse = torch.max(torch.log1p(torch.relu(logits)), dim=1).values.squeeze()
    # Return non-zero indices and values as a sparse dict
    nz_indices = sparse.nonzero().squeeze().tolist()
    nz_values = sparse[nz_indices].tolist()
    return {'indices': nz_indices, 'values': nz_values}

def upsert_hybrid_docs(index, documents: list[dict], batch_size: int = 100):
    records = []
    for doc in documents:
        dense_vec = embed(doc['text'])        # OpenAI embedding
        sparse_vec = generate_sparse_vector(doc['text'])  # SPLADE
        records.append({
            'id': doc['id'],
            'values': dense_vec,
            'sparse_values': sparse_vec,
            'metadata': {'text': doc['text'], 'source': doc['source']},
        })

    for i in range(0, len(records), batch_size):
        batch = records[i:i + batch_size]
        index.upsert(vectors=batch)
        print(f'Upserted batch {i//batch_size + 1}')

def hybrid_query_pinecone(index, query: str, alpha: float = 0.5, top_k: int = 5):
    dense_vec = embed(query)
    sparse_vec = generate_sparse_vector(query)

    # Scale vectors by alpha for weighted fusion
    scaled_dense = [v * alpha for v in dense_vec]
    scaled_sparse = {
        'indices': sparse_vec['indices'],
        'values': [v * (1 - alpha) for v in sparse_vec['values']],
    }

    results = index.query(
        vector=scaled_dense,
        sparse_vector=scaled_sparse,
        top_k=top_k,
        include_metadata=True,
    )
    return results.matches

-- Setup: enable extensions
CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pg_trgm;

-- Create hybrid-capable table
CREATE TABLE documents (
    id SERIAL PRIMARY KEY,
    content TEXT NOT NULL,
    embedding vector(1536),
    content_tsv tsvector GENERATED ALWAYS AS (to_tsvector('english', content)) STORED
);

-- Indexes for both search modes
CREATE INDEX ON documents USING hnsw (embedding vector_cosine_ops);
CREATE INDEX ON documents USING GIN (content_tsv);

import psycopg2
import json

def pgvector_hybrid_search(conn, query: str, top_k: int = 5):
    dense_vec = embed(query)
    tsquery = ' & '.join(query.split())  # simple AND query

    # Dense query
    dense_sql = '''
        SELECT id, content, 1 - (embedding <=> %s::vector) AS score
        FROM documents
        ORDER BY embedding <=> %s::vector
        LIMIT 20
    '''
    # Sparse query
    sparse_sql = '''
        SELECT id, content, ts_rank_cd(content_tsv, to_tsquery('english', %s)) AS score
        FROM documents
        WHERE content_tsv @@ to_tsquery('english', %s)
        ORDER BY score DESC
        LIMIT 20
    '''
    with conn.cursor() as cur:
        cur.execute(dense_sql, [json.dumps(dense_vec), json.dumps(dense_vec)])
        dense_rows = cur.fetchall()
        cur.execute(sparse_sql, [tsquery, tsquery])
        sparse_rows = cur.fetchall()

    return fuse_with_rrf(dense_rows, sparse_rows, top_k)

def fuse_with_rrf(dense_rows, sparse_rows, top_k: int = 5, k: int = 60):
    from collections import defaultdict

    doc_texts = {}
    scores = defaultdict(float)

    for rank, row in enumerate(dense_rows, start=1):
        doc_id, text, _ = row
        scores[doc_id] += 1.0 / (k + rank)
        doc_texts[doc_id] = text

    for rank, row in enumerate(sparse_rows, start=1):
        doc_id, text, _ = row
        scores[doc_id] += 1.0 / (k + rank)
        doc_texts[doc_id] = text

    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return [{'id': doc_id, 'text': doc_texts[doc_id], 'rrf_score': s}
            for doc_id, s in ranked[:top_k]]

def evaluate_retrieval(search_fn, test_queries, golden_relevant, k=5):
    ndcg_scores = []
    hit_count = 0

    for query, relevant_ids in zip(test_queries, golden_relevant):
        results = search_fn(query, top_k=k)
        retrieved_ids = [r['id'] for r in results]

        # Hit rate
        if any(rid in relevant_ids for rid in retrieved_ids):
            hit_count += 1

        # Simplified NDCG (binary relevance)
        dcg = sum(
            1.0 / (i + 1)
            for i, rid in enumerate(retrieved_ids)
            if rid in relevant_ids
        )
        idcg = sum(1.0 / (i + 1) for i in range(min(len(relevant_ids), k)))
        ndcg_scores.append(dcg / idcg if idcg > 0 else 0)

    return {
        'hit_rate': hit_count / len(test_queries),
        'ndcg': sum(ndcg_scores) / len(ndcg_scores),
    }

def find_optimal_alpha(index, test_queries, golden_relevant, alphas=None):
    if alphas is None:
        alphas = [0.0, 0.2, 0.4, 0.5, 0.6, 0.8, 1.0]

    best_alpha, best_score = 0.5, 0
    for alpha in alphas:
        hits = 0
        for query, relevant in zip(test_queries, golden_relevant):
            results = hybrid_query_pinecone(index, query, alpha=alpha, top_k=5)
            if any(r['id'] in relevant for r in results):
                hits += 1
        hit_rate = hits / len(test_queries)
        print(f'alpha={alpha}: hit_rate={hit_rate:.3f}')
        if hit_rate > best_score:
            best_score, best_alpha = hit_rate, alpha
    return best_alpha