First push

.gitignore

@@ -0,0 +1,2 @@
+.venv
+__pycache__

README.md

@@ -0,0 +1,15 @@
+---
+title: ZSInvert
+emoji: 🔄
+colorFrom: blue
+colorTo: green
+sdk: gradio
+sdk_version: 6.5.1
+python_version: "3.10"
+app_file: app.py
+pinned: false
+license: mit
+short_description: Reconstruct text from embedding vectors via beam search
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,280 @@
+"""
+ZSInvert — Zero-Shot Embedding Inversion Explorer.
+
+Interactive tool demonstrating embedding inversion via
+adversarial decoding beam search. Reconstructs text from
+embedding vectors without training embedding-specific models.
+
+Part of E04: ZSInvert.
+"""
+
+import time
+import threading
+import queue
+
+import gradio as gr
+import torch
+
+try:
+    import spaces
+    gpu_decorator = spaces.GPU(duration=120)
+except ImportError:
+    gpu_decorator = lambda fn: fn
+
+from model import load_llm, load_encoder, encode_text, ENCODERS
+from invert import beam_search
+
+_STAGE1_PROMPT = "tell me a story"
+_STAGE2_PROMPT_TEMPLATE = "write a sentence similar to this: {seed}"
+
+# Encoder choices (drop contriever — broken)
+_ENCODER_CHOICES = [k for k in ENCODERS if k != "contriever"]
+
+
+def _sim_color(cos_sim: float) -> str:
+    """Return hex color for a cosine similarity value."""
+    if cos_sim > 0.99:
+        return "#3b82f6"  # blue
+    if cos_sim > 0.95:
+        return "#16a34a"  # dark green
+    if cos_sim > 0.85:
+        return "#65a30d"  # green
+    if cos_sim > 0.70:
+        return "#ca8a04"  # amber
+    if cos_sim > 0.50:
+        return "#ef4444"  # red
+    return "#a855f7"      # purple
+
+
+def _format_results(stage_results: list[dict]) -> str:
+    """Render accumulated stage results as styled HTML."""
+    if not stage_results:
+        return ""
+    rows = []
+    for r in stage_results:
+        color = _sim_color(r["cos_sim"])
+        rows.append(
+            f'<div style="margin-bottom:12px;padding:10px;border:1px solid #333;border-radius:6px;'
+            f'background:#1a1a2e;">'
+            f'<span style="font-weight:bold;color:#ccc;">S{r["stage"]}</span> '
+            f'<span style="color:#eee;font-style:italic;">"{r["text"]}"</span><br>'
+            f'<span style="color:{color};font-weight:bold;">cos={r["cos_sim"]:.4f}</span>'
+            f'&nbsp;&nbsp;len={r["length"]}'
+            f'&nbsp;&nbsp;{r["time"]:.1f}s'
+            f'&nbsp;&nbsp;steps={r["steps"]}'
+            f'</div>'
+        )
+    return "".join(rows)
+
+
+def _format_progress(step: int, text: str, cos_sim: float) -> str:
+    """Render a live progress line during beam search."""
+    color = _sim_color(cos_sim)
+    return (
+        f'<div style="padding:8px;border:1px dashed #555;border-radius:4px;'
+        f'background:#111;margin-bottom:12px;">'
+        f'<span style="color:#888;">step {step}</span>&nbsp;&nbsp;'
+        f'<span style="color:{color};font-weight:bold;">cos={cos_sim:.4f}</span>&nbsp;&nbsp;'
+        f'<span style="color:#aaa;font-style:italic;">"{text}"</span>'
+        f'</div>'
+    )
+
+
+_SENTINEL = object()
+
+
+@gpu_decorator
+def _run_beam_search_threaded(
+    target_emb, encoder_name, prompt,
+    beam_width, top_k, patience, max_steps, min_similarity, randomness,
+    progress_queue,
+):
+    """Run beam search on GPU, pushing step updates to a queue."""
+    llm, tokenizer = load_llm()
+    encoder = load_encoder(encoder_name)
+
+    step_count = 0
+
+    def on_step(step, cand):
+        nonlocal step_count
+        step_count = step
+        progress_queue.put((step, cand.seq_str, cand.cos_sim))
+
+    t0 = time.time()
+    result = beam_search(
+        llm, tokenizer, encoder, target_emb,
+        prompt=prompt,
+        beam_width=int(beam_width),
+        max_steps=int(max_steps),
+        top_k=int(top_k),
+        patience=int(patience),
+        min_similarity=float(min_similarity),
+        randomness=bool(randomness),
+        on_step=on_step,
+    )
+    elapsed = time.time() - t0
+    progress_queue.put(_SENTINEL)
+    return result, elapsed, step_count
+
+
+def run_stage(
+    text, encoder_name,
+    beam_width, top_k, patience, max_steps, min_similarity, randomness,
+    target_emb_state, stage_results_state,
+):
+    """Run the next stage of inversion, yielding progress updates."""
+    if not text or not text.strip():
+        gr.Warning("Please enter some text.")
+        yield (
+            target_emb_state,
+            stage_results_state,
+            _format_results(stage_results_state),
+            gr.update(),
+        )
+        return
+
+    stage_num = len(stage_results_state) + 1
+
+    # Encode target on first stage
+    if stage_num == 1:
+        encoder = load_encoder(encoder_name)
+        target_emb_state = encode_text(text.strip(), encoder)
+
+    # Build prompt
+    if stage_num == 1:
+        prompt = _STAGE1_PROMPT
+    else:
+        prev_text = stage_results_state[-1]["text"]
+        prompt = _STAGE2_PROMPT_TEMPLATE.format(seed=prev_text)
+
+    # Run beam search in a thread so we can yield progress
+    progress_q = queue.Queue()
+
+    # Container for the thread's return value
+    result_holder = [None, 0.0, 0]
+
+    def _worker():
+        r, elapsed, steps = _run_beam_search_threaded(
+            target_emb_state, encoder_name, prompt,
+            beam_width, top_k, patience, max_steps, min_similarity, randomness,
+            progress_q,
+        )
+        result_holder[0] = r
+        result_holder[1] = elapsed
+        result_holder[2] = steps
+
+    worker = threading.Thread(target=_worker)
+    worker.start()
+
+    # Yield progress updates as they arrive
+    completed_html = _format_results(stage_results_state)
+    while True:
+        try:
+            item = progress_q.get(timeout=0.5)
+        except queue.Empty:
+            if not worker.is_alive():
+                break
+            continue
+
+        if item is _SENTINEL:
+            break
+
+        step, best_text, best_sim = item
+        progress_html = _format_progress(step, best_text, best_sim)
+        yield (
+            target_emb_state,
+            stage_results_state,
+            completed_html + progress_html,
+            gr.update(value=f"Running S{stage_num}...", interactive=False),
+        )
+
+    worker.join()
+
+    result, elapsed, steps = result_holder
+    stage_results_state = stage_results_state + [{
+        "stage": stage_num,
+        "text": result.seq_str,
+        "cos_sim": result.cos_sim,
+        "length": len(result.token_ids),
+        "time": elapsed,
+        "steps": steps,
+    }]
+
+    html = _format_results(stage_results_state)
+    btn_label = f"Run Stage {stage_num + 1}"
+
+    yield (
+        target_emb_state,
+        stage_results_state,
+        html,
+        gr.update(value=btn_label, visible=True, interactive=True),
+    )
+
+
+def reset_state():
+    """Reset all state for a fresh run."""
+    return None, [], "", gr.update(value="Run Stage 1", visible=True, interactive=True)
+
+
+with gr.Blocks(title="ZSInvert") as demo:
+    gr.Markdown("# ZSInvert — Zero-Shot Embedding Inversion")
+    gr.Markdown(
+        "Reconstruct text from its embedding vector using "
+        "cosine-similarity-guided beam search."
+    )
+
+    # --- State ---
+    target_emb_state = gr.State(value=None)
+    stage_results_state = gr.State(value=[])
+
+    # --- Input row ---
+    with gr.Row():
+        text_input = gr.Textbox(
+            label="Input text",
+            placeholder="Enter text to encode and invert...",
+            scale=4,
+        )
+        encoder_dd = gr.Dropdown(
+            choices=_ENCODER_CHOICES,
+            value="gte",
+            label="Encoder",
+            scale=1,
+        )
+
+    # --- Advanced settings ---
+    with gr.Accordion("Advanced Settings", open=False):
+        with gr.Row():
+            beam_width_sl = gr.Slider(5, 50, value=10, step=1, label="beam_width")
+            top_k_sl = gr.Slider(5, 50, value=10, step=1, label="top_k")
+            patience_sl = gr.Slider(0, 20, value=5, step=1, label="patience (0=off)")
+        with gr.Row():
+            max_steps_sl = gr.Slider(0, 64, value=0, step=1, label="max_steps (0=unlimited)")
+            min_sim_sl = gr.Slider(0.0, 1.0, value=0.0, step=0.01, label="min_similarity (0=off)")
+            randomness_cb = gr.Checkbox(value=True, label="randomness")
+
+    # --- Run button ---
+    run_btn = gr.Button("Run Stage 1", variant="primary")
+
+    # --- Results ---
+    results_html = gr.HTML(value="", label="Results")
+
+    # --- Wiring ---
+    all_inputs = [
+        text_input, encoder_dd,
+        beam_width_sl, top_k_sl, patience_sl, max_steps_sl, min_sim_sl, randomness_cb,
+        target_emb_state, stage_results_state,
+    ]
+    all_outputs = [
+        target_emb_state, stage_results_state,
+        results_html, run_btn,
+    ]
+
+    run_btn.click(fn=run_stage, inputs=all_inputs, outputs=all_outputs)
+
+    # Reset when input text or encoder changes
+    text_input.change(fn=reset_state, inputs=[], outputs=all_outputs)
+    encoder_dd.change(fn=reset_state, inputs=[], outputs=all_outputs)
+
+
+if __name__ == "__main__":
+    demo.launch(server_port=7860, theme=gr.themes.Base())

invert.py

@@ -0,0 +1,412 @@
+"""
+Beam search inversion engine for ZSInvert.
+
+Cosine-similarity-guided beam search that reconstructs text
+from an embedding vector using a small LLM as the token
+proposal engine.
+
+Part of E04: ZSInvert.
+"""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass, field
+from typing import Callable
+
+import torch
+import torch.nn.functional as F
+from sentence_transformers import SentenceTransformer
+from transformers import AutoModelForCausalLM, AutoTokenizer, DynamicCache
+
+from model import get_chat_format
+
+# Tokens to mask from generation (special/formatting tokens)
+_MASK_STRINGS = [
+    "<|im_end|>", "<|end_header_id|>", "<|start_header_id|>",
+    "<|eot_id|>", "<|eom_id|>", "<|python_tag|>",
+    "@", "\xa0", '"', "\n", "\n\n", " \n\n",
+]
+
+# Number of top beams kept deterministically in randomness mode
+_FIXED_KEEP = 5
+
+
+@dataclass
+class Candidate:
+    """A beam search candidate."""
+    token_ids: list[int] = field(default_factory=list)
+    seq_str: str = ""
+    score: float = 0.0
+    cos_sim: float = 0.0
+    kv_cache: DynamicCache | None = field(default=None, repr=False)
+
+
+@dataclass
+class InversionResult:
+    """Result of a full inversion run."""
+    original_text: str | None = None
+    target_embedding: torch.Tensor | None = None
+    stage1_text: str = ""
+    stage1_cos_sim: float = 0.0
+    stage2_text: str = ""
+    stage2_cos_sim: float = 0.0
+
+
+def _top_k_top_p_filter(logits: torch.Tensor, top_k: int, top_p: float) -> list[int]:
+    """Return indices that survive top-k and top-p filtering."""
+    # Top-k: keep only top_k highest logits
+    topk_vals, topk_idx = torch.topk(logits, min(top_k, logits.size(-1)))
+
+    # Top-p (nucleus): keep smallest set whose cumulative prob >= top_p
+    probs = F.softmax(topk_vals, dim=-1)
+    cumulative = torch.cumsum(probs, dim=-1)
+    # Mask tokens beyond the nucleus
+    mask = cumulative - probs <= top_p
+    filtered_idx = topk_idx[mask]
+
+    return filtered_idx.tolist()
+
+
+_cached_mask_ids: list[int] | None = None
+
+
+def _build_mask_token_ids(tokenizer: AutoTokenizer) -> list[int]:
+    """Build set of token IDs to suppress during generation. Cached.
+
+    Masks both exact single-token matches for _MASK_STRINGS and any
+    vocab token whose decoded form contains a newline (catches merged
+    tokens like '.\\n' that bypass the single-token check).
+    """
+    global _cached_mask_ids
+    if _cached_mask_ids is not None:
+        return _cached_mask_ids
+
+    mask_ids = set()
+    for s in _MASK_STRINGS:
+        tokens = tokenizer.encode(s, add_special_tokens=False)
+        if len(tokens) == 1:
+            mask_ids.add(tokens[0])
+    if tokenizer.eos_token_id is not None:
+        mask_ids.add(tokenizer.eos_token_id)
+    # Also mask any vocab token containing a newline
+    for tid in range(tokenizer.vocab_size):
+        decoded = tokenizer.decode([tid])
+        if "\n" in decoded:
+            mask_ids.add(tid)
+    _cached_mask_ids = list(mask_ids)
+    return _cached_mask_ids
+
+
+def _get_next_token_candidates(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    prefix: list[int],
+    suffix: list[int],
+    prompt_tokens: list[int],
+    candidates: list[Candidate],
+    top_k: int,
+    top_p: float,
+    repetition_penalty: float,
+    mask_ids: list[int],
+) -> list[list[tuple[int, float]]]:
+    """Forward pass through LLM to get candidate next tokens.
+
+    Builds input as: prefix + prompt_tokens + suffix + candidate.token_ids
+    Uses KV-cache from candidates when available.
+
+    Returns list of [(token_id, log_prob), ...] per candidate.
+    """
+    device = next(model.parameters()).device
+
+    # Build full token sequences
+    base = prefix + prompt_tokens + suffix
+    batch_tokens = [base + c.token_ids for c in candidates]
+
+    # All sequences should have the same length (beam search invariant)
+    assert len(set(len(t) for t in batch_tokens)) == 1
+
+    input_ids = torch.tensor(batch_tokens, device=device)
+
+    # Check for usable KV-cache
+    batch_kv = [c.kv_cache for c in candidates]
+    use_cache = all(kv is not None for kv in batch_kv)
+
+    if use_cache:
+        kv_cache = DynamicCache.from_batch_splits(batch_kv)
+        cache_len = kv_cache.get_seq_length()
+        model_input = input_ids[:, cache_len:]
+        attn_mask = torch.ones_like(input_ids, device=device)
+    else:
+        kv_cache = DynamicCache()
+        model_input = input_ids
+        attn_mask = None
+
+    with torch.no_grad():
+        outputs = model(
+            input_ids=model_input,
+            attention_mask=attn_mask,
+            past_key_values=kv_cache,
+            use_cache=True,
+        )
+
+    # Split KV-cache back per candidate
+    next_kv = outputs.past_key_values
+    try:
+        split_kv = next_kv.batch_split(len(candidates), 1) if next_kv else [None] * len(candidates)
+    except Exception:
+        split_kv = [None] * len(candidates)
+
+    logits = outputs.logits[:, -1, :]  # (batch, vocab)
+
+    # Apply repetition penalty
+    if repetition_penalty != 1.0:
+        for i, tokens in enumerate(batch_tokens):
+            for tid in set(tokens):
+                if logits[i, tid] > 0:
+                    logits[i, tid] /= repetition_penalty
+                else:
+                    logits[i, tid] *= repetition_penalty
+
+    # Mask special tokens
+    logits[:, mask_ids] = -1e10
+
+    log_probs = F.log_softmax(logits, dim=-1)
+
+    results = []
+    for i in range(len(candidates)):
+        filtered = _top_k_top_p_filter(logits[i], top_k, top_p)
+        pairs = [(tid, log_probs[i, tid].item()) for tid in filtered]
+        pairs.sort(key=lambda x: x[1], reverse=True)
+        results.append(pairs)
+
+    return results, split_kv
+
+
+def _score_candidates(
+    encoder: SentenceTransformer,
+    target_embedding: torch.Tensor,
+    candidates: list[Candidate],
+) -> None:
+    """Score candidates by cosine similarity to target embedding. Mutates in place."""
+    if not candidates:
+        return
+
+    texts = [c.seq_str for c in candidates]
+    embs = encoder.encode(texts, convert_to_tensor=True, normalize_embeddings=True)
+
+    # target_embedding shape: (1, dim) — broadcast
+    target_norm = F.normalize(target_embedding, dim=-1)
+    sims = torch.matmul(embs, target_norm.squeeze(0))  # (batch,)
+
+    for i, c in enumerate(candidates):
+        c.cos_sim = sims[i].item()
+        c.score = c.cos_sim
+
+
+def beam_search(
+    model: AutoModelForCausalLM,
+    tokenizer: AutoTokenizer,
+    encoder: SentenceTransformer,
+    target_embedding: torch.Tensor,
+    prompt: str,
+    beam_width: int = 30,
+    max_steps: int = 0,
+    top_k: int = 30,
+    top_p: float = 1.0,
+    repetition_penalty: float = 1.5,
+    randomness: bool = True,
+    patience: int = 5,
+    min_similarity: float = 0.0,
+    on_step: Callable | None = None,
+) -> Candidate:
+    """Run cosine-similarity-guided beam search.
+
+    Args:
+        model: Generator LLM.
+        tokenizer: LLM tokenizer.
+        encoder: Embedding encoder for scoring.
+        target_embedding: Target embedding to invert. Shape (1, dim).
+        prompt: User-facing prompt (becomes chat user message).
+        beam_width: Number of candidates to maintain per step.
+        max_steps: Maximum tokens to generate. 0 means no limit (stop via patience only).
+        top_k: Top-k tokens to consider per expansion.
+        top_p: Nucleus sampling threshold.
+        repetition_penalty: Penalty for repeated tokens in logits.
+        randomness: If True, keep top 5 deterministically + sample rest.
+        patience: Stop after this many steps with no improvement in best cosine sim.
+            Set to 0 to disable early stopping.
+        min_similarity: Stop immediately when cosine sim reaches this threshold.
+            Set to 0.0 to disable.
+        on_step: Callback(step, best_candidate) fired each step.
+
+    Returns:
+        Best candidate found during search.
+    """
+    prefix, suffix = get_chat_format(tokenizer)
+    prompt_tokens = tokenizer.encode(prompt, add_special_tokens=False)
+    mask_ids = _build_mask_token_ids(tokenizer)
+
+    candidates = [Candidate()]
+    best_complete: Candidate | None = None
+    best_ever: Candidate | None = None
+    steps_since_improvement = 0
+
+    step = 0
+    while max_steps <= 0 or step < max_steps:
+        step += 1
+        # Expand: get next-token proposals for each candidate
+        token_proposals, split_kv = _get_next_token_candidates(
+            model, tokenizer, prefix, suffix, prompt_tokens,
+            candidates, top_k, top_p, repetition_penalty, mask_ids,
+        )
+
+        # Build expanded candidates
+        expanded: list[Candidate] = []
+        for i, cand in enumerate(candidates):
+            for tid, _logp in token_proposals[i]:
+                new_ids = cand.token_ids + [tid]
+                expanded.append(Candidate(
+                    token_ids=new_ids,
+                    seq_str=tokenizer.decode(new_ids),
+                    kv_cache=split_kv[i] if split_kv[i] is not None else None,
+                ))
+
+        # Score by cosine similarity
+        _score_candidates(encoder, target_embedding, expanded)
+
+        # Sort by score descending
+        expanded.sort(key=lambda c: c.score, reverse=True)
+
+        # Track best-ever candidate (highest cosine sim at any step)
+        step_best = expanded[0]
+        if best_ever is None or step_best.cos_sim > best_ever.cos_sim:
+            best_ever = Candidate(
+                token_ids=list(step_best.token_ids),
+                seq_str=step_best.seq_str,
+                score=step_best.score,
+                cos_sim=step_best.cos_sim,
+            )
+            steps_since_improvement = 0
+        else:
+            steps_since_improvement += 1
+            if patience > 0 and steps_since_improvement >= patience:
+                break
+
+        if min_similarity > 0 and best_ever.cos_sim >= min_similarity:
+            break
+
+        # Track best complete sentence
+        for c in expanded:
+            if c.seq_str and c.seq_str.rstrip()[-1:] in ".?!":
+                if best_complete is None or c.score > best_complete.score:
+                    best_complete = Candidate(
+                        token_ids=list(c.token_ids),
+                        seq_str=c.seq_str,
+                        score=c.score,
+                        cos_sim=c.cos_sim,
+                    )
+
+        # Select: top beam_width candidates (with optional randomness)
+        if randomness and len(expanded) > _FIXED_KEEP:
+            keep = min(_FIXED_KEEP, beam_width)
+            remainder = min(beam_width - keep, len(expanded) - keep)
+            candidates = expanded[:keep]
+            if remainder > 0:
+                candidates += random.sample(expanded[keep:], remainder)
+        else:
+            candidates = expanded[:beam_width]
+
+        # Callback
+        if on_step is not None:
+            best_so_far = best_complete if best_complete else candidates[0]
+            on_step(step, best_so_far)
+
+    # Return the candidate with the highest cosine similarity across all tracking
+    finalists = [c for c in [best_ever, best_complete, candidates[0]] if c is not None]
+    return max(finalists, key=lambda c: c.cos_sim)
+
+
+_STAGE1_PROMPT = "tell me a story"
+_STAGE2_PROMPT_TEMPLATE = "write a sentence similar to this: {seed}"
+
+
+def invert(
+    text: str,
+    encoder_name: str = "gte",
+    beam_width: int = 30,
+    max_steps: int = 0,
+    top_k: int = 30,
+    two_stage: bool = True,
+    on_progress: Callable | None = None,
+) -> InversionResult:
+    """Run the full two-stage ZSInvert inversion pipeline.
+
+    Stage 1: Seed generation with a generic prompt.
+    Stage 2: Paraphrase refinement using the Stage 1 output as context.
+
+    Args:
+        text: Input text to encode and then invert.
+        encoder_name: Which embedding encoder to use ("gte", "gtr", "contriever").
+        beam_width: Beam search width.
+        max_steps: Maximum tokens per stage.
+        top_k: Top-k tokens per expansion step.
+        two_stage: If True, run both stages. If False, Stage 1 only.
+        on_progress: Callback(stage, step, best_candidate) for UI updates.
+            stage is 1 or 2, step is the beam search step index.
+
+    Returns:
+        InversionResult with results from both stages.
+    """
+    from model import load_llm, load_encoder, encode_text
+
+    model, tokenizer = load_llm()
+    encoder = load_encoder(encoder_name)
+    target_embedding = encode_text(text, encoder)
+
+    # Stage 1: seed generation
+    def stage1_callback(step: int, cand: Candidate) -> None:
+        if on_progress is not None:
+            on_progress(1, step, cand)
+
+    stage1 = beam_search(
+        model, tokenizer, encoder, target_embedding,
+        prompt=_STAGE1_PROMPT,
+        beam_width=beam_width,
+        max_steps=max_steps,
+        top_k=top_k,
+        randomness=True,
+        on_step=stage1_callback,
+    )
+
+    result = InversionResult(
+        original_text=text,
+        target_embedding=target_embedding,
+        stage1_text=stage1.seq_str,
+        stage1_cos_sim=stage1.cos_sim,
+    )
+
+    if not two_stage:
+        result.stage2_text = result.stage1_text
+        result.stage2_cos_sim = result.stage1_cos_sim
+        return result
+
+    # Stage 2: paraphrase refinement
+    def stage2_callback(step: int, cand: Candidate) -> None:
+        if on_progress is not None:
+            on_progress(2, step, cand)
+
+    stage2_prompt = _STAGE2_PROMPT_TEMPLATE.format(seed=stage1.seq_str)
+    stage2 = beam_search(
+        model, tokenizer, encoder, target_embedding,
+        prompt=stage2_prompt,
+        beam_width=beam_width,
+        max_steps=max_steps,
+        top_k=top_k,
+        randomness=True,
+        on_step=stage2_callback,
+    )
+
+    result.stage2_text = stage2.seq_str
+    result.stage2_cos_sim = stage2.cos_sim
+    return result

model.py

@@ -0,0 +1,103 @@
+"""
+Model loading for ZSInvert.
+
+Loads the generator LLM (Qwen2.5-0.5B-Instruct) and selectable
+embedding encoders (GTE-base, GTR-T5-base, Contriever).
+
+Part of E04: ZSInvert.
+"""
+
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from sentence_transformers import SentenceTransformer
+
+GENERATOR_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
+
+ENCODERS = {
+    "gte": "thenlper/gte-base",
+    "gtr": "sentence-transformers/gtr-t5-base",
+    "contriever": "facebook/contriever",
+    "mini": "sentence-transformers/all-MiniLM-L6-v2",
+}
+
+_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+_llm: AutoModelForCausalLM | None = None
+_llm_tokenizer: AutoTokenizer | None = None
+_encoders: dict[str, SentenceTransformer] = {}
+
+
+def load_llm() -> tuple[AutoModelForCausalLM, AutoTokenizer]:
+    """Load generator LLM. Singleton."""
+    global _llm, _llm_tokenizer
+    if _llm is None:
+        _llm_tokenizer = AutoTokenizer.from_pretrained(GENERATOR_MODEL)
+        _llm = AutoModelForCausalLM.from_pretrained(
+            GENERATOR_MODEL,
+            dtype=torch.bfloat16,
+        ).eval().to(_device)
+    return _llm, _llm_tokenizer
+
+
+def load_encoder(name: str = "gte") -> SentenceTransformer:
+    """Load embedding encoder by name. Cached per name."""
+    if name not in ENCODERS:
+        raise ValueError(f"Unknown encoder '{name}'. Choose from: {list(ENCODERS.keys())}")
+    if name not in _encoders:
+        model_id = ENCODERS[name]
+        _encoders[name] = SentenceTransformer(model_id, device=_device)
+    return _encoders[name]
+
+
+def encode_text(text: str, encoder: SentenceTransformer) -> torch.Tensor:
+    """Encode text to normalized embedding vector. Returns shape (1, hidden_dim)."""
+    emb = encoder.encode(
+        text,
+        convert_to_tensor=True,
+        normalize_embeddings=True,
+    )
+    return emb.unsqueeze(0)
+
+
+def get_chat_format(tokenizer: AutoTokenizer) -> tuple[list[int], list[int]]:
+    """Extract chat prefix/suffix token IDs from the Qwen2.5 chat template.
+
+    The prefix is everything the template adds before the user content.
+    The suffix is everything after the user content through the generation prompt.
+
+    For Qwen2.5 the structure is:
+        <|im_start|>system\\n...system prompt...<|im_end|>\\n
+        <|im_start|>user\\n{CONTENT}<|im_end|>\\n
+        <|im_start|>assistant\\n
+
+    We split so that: prefix + prompt_tokens + suffix = full template.
+    """
+    # Template with empty content (no gen prompt) — find where content is inserted
+    empty = tokenizer.apply_chat_template(
+        [{"role": "user", "content": ""}],
+        add_generation_prompt=False,
+    )
+    # Template with a known marker to locate the split point
+    marker = tokenizer.apply_chat_template(
+        [{"role": "user", "content": "hello"}],
+        add_generation_prompt=True,
+    )
+    marker_tokens = tokenizer.encode("hello", add_special_tokens=False)
+
+    # Find where the marker content appears in the full template
+    marker_len = len(marker_tokens)
+    for i in range(len(marker)):
+        if marker[i : i + marker_len] == marker_tokens:
+            prefix = marker[:i]
+            suffix = marker[i + marker_len :]
+            return prefix, suffix
+
+    # Fallback: use the empty template structure
+    # Empty template has <|im_end|>\n right after user\n — drop those 2 tokens
+    prefix = empty[:-2]
+    full_gen = tokenizer.apply_chat_template(
+        [{"role": "user", "content": ""}],
+        add_generation_prompt=True,
+    )
+    suffix = full_gen[len(prefix):]
+    return prefix, suffix

requirements.txt

@@ -0,0 +1 @@
+nltk