"""baa-ai-Embedding-Reranker-v1 — unified embedder+reranker over a shared word-embedding table.
The reranker's word-embedding matrix is stored once (in the embedder) and tied at load => ~23% footprint."""
import os, glob, torch
import torch.nn.functional as F
from safetensors.torch import load_file
from sentence_transformers import SentenceTransformer
from transformers import AutoConfig, AutoModelForSequenceClassification, AutoTokenizer

class BaaEmbeddingReranker:
    def __init__(self, path=None, device=None):
        path = path or os.path.dirname(os.path.abspath(__file__))
        self.device = device or ("mps" if torch.backends.mps.is_available() else "cpu")
        emb_dir, rr_dir = os.path.join(path, "embedder"), os.path.join(path, "reranker")
        # embedder = bi-encoder stack (holds the canonical shared word-embedding table)
        self.embedder = SentenceTransformer(emb_dir, trust_remote_code=True, device=self.device)
        shared_wemb = self.embedder[0].auto_model.embeddings.word_embeddings.weight.data
        # reranker = cross-encoder seq-classifier, word-emb injected from the shared table (stripped on disk)
        cfg = AutoConfig.from_pretrained(rr_dir)
        self.reranker = AutoModelForSequenceClassification.from_config(cfg).half()
        sf = glob.glob(os.path.join(rr_dir, "**", "*.safetensors"), recursive=True)[0]
        self.reranker.load_state_dict(load_file(sf), strict=False)      # word-emb missing -> injected next
        self.reranker.roberta.embeddings.word_embeddings.weight.data = shared_wemb.to(self.reranker.dtype)
        self.reranker.to(self.device).eval()
        self.rr_tok = AutoTokenizer.from_pretrained(rr_dir)

    def embed(self, texts, is_query=False, batch_size=32):
        pref = "query: " if is_query else ""
        return self.embedder.encode([pref + t for t in texts], normalize_embeddings=True,
                                    batch_size=batch_size, show_progress_bar=False)

    @torch.no_grad()
    def rerank(self, query, docs, top_k=None, batch_size=32):
        scores = []
        for i in range(0, len(docs), batch_size):
            enc = self.rr_tok([(query, d[:2000]) for d in docs[i:i+batch_size]], padding=True,
                              truncation=True, max_length=512, return_tensors="pt").to(self.device)
            scores.extend(self.reranker(**enc).logits[:, 0].float().cpu().tolist())
        order = sorted(range(len(docs)), key=lambda j: -scores[j])
        if top_k:
            order = order[:top_k]
        return [(docs[j], scores[j]) for j in order]

if __name__ == "__main__":
    m = BaaEmbeddingReranker()
    q = "How does a cross-encoder reranker work?"
    docs = ["A cross-encoder jointly encodes the query and document to score relevance.",
            "The mitochondria is the powerhouse of the cell.",
            "Bi-encoders embed query and document separately for fast retrieval."]
    qv = m.embed([q], is_query=True)[0]; dv = m.embed(docs)
    import numpy as np
    print("embed cos:", [round(float(np.dot(qv, d)), 3) for d in dv])
    print("rerank   :", [(round(s, 2), d[:45]) for d, s in m.rerank(q, docs)])