import pandas as pd
import re

def clean_model_name(name: str) -> str:
    """Strip markdown link syntax: [ModelName](url) -> ModelName"""
    match = re.search(r'\[(.+?)\]', name)
    return match.group(1) if match else name

summary = pd.read_csv("mteb_lb_ds_code_feb_25_2026_summary.csv")
tasks   = pd.read_csv("mteb_lb_ds_code_feb_25_2026_perf_per_task.csv")

summary["Model"] = summary["Model"].apply(clean_model_name)

print("Task columns in perf_per_task:")
print([c for c in tasks.columns if c != "Model"])

Task columns in perf_per_task:
['AppsRetrieval', 'COIRCodeSearchNetRetrieval', 'CodeEditSearchRetrieval',
 'CodeFeedbackMT', 'CodeFeedbackST', 'CodeSearchNetCCRetrieval',
 'CodeSearchNetRetrieval', 'CodeTransOceanContest', 'CodeTransOceanDL',
 'CosQA', 'StackOverflowQA', 'SyntheticText2SQL']

merged = summary.merge(tasks, on="Model", how="inner")

print(merged["Retrieval"].describe())
print(f"\nModels with retrieval >= 90: {(merged['Retrieval'] >= 90).sum()}")
print(f"Models with retrieval >= 80: {(merged['Retrieval'] >= 80).sum()}")
print(f"Models with retrieval < 70:  {(merged['Retrieval'] < 70).sum()}")

count    204.000000
mean      64.318627
std       17.421083
min       19.220000
25%       54.417500
50%       66.185000
75%       77.320000
max       97.290000

Models with retrieval >= 90: 9
Models with retrieval >= 80: 34
Models with retrieval < 70:  108

import matplotlib.pyplot as plt

local = merged.dropna(subset=["Memory Usage (MB)", "Retrieval"]).copy()
local["mem_gb"] = local["Memory Usage (MB)"] / 1024

highlights = {
    "CodeSearch-ModernBERT-Crow-Plus": (0.59, 84.59),
    "Octen-Embedding-4B":              (7.49, 90.81),
    "Octen-Embedding-8B":              (14.10, 93.36),
    "llama-embed-nemotron-8b":         (27.96, 96.59),
    "Qwen3-Embedding-4B":              (7.49, 80.07),
}

fig, ax = plt.subplots(figsize=(10, 6))
ax.scatter(local["mem_gb"], local["Retrieval"], alpha=0.5, s=40, color="#aaa")

for name, (x, y) in highlights.items():
    ax.scatter(x, y, s=80, zorder=5)
    ax.annotate(
        name,
        xy=(x, y),
        xytext=(x + 0.4, y - 1.8),
        fontsize=8,
        arrowprops=dict(arrowstyle="->", color="#555", lw=0.8),
    )

ax.set_xlabel("VRAM / RAM (GB)")
ax.set_ylabel("Retrieval Score")
ax.set_title("MTEB Code — Retrieval Score vs Memory (self-hosted only)")
plt.tight_layout()
plt.savefig("retrieval_vs_memory.png", dpi=150)

import voyageai

client = voyageai.Client()  # reads VOYAGE_API_KEY from env

def embed_contracts(
    paths: list[str],
    model: str = "voyage-4-large",
    batch_size: int = 128,
) -> list[list[float]]:
    """
    Embed Solidity files by path.
    Voyage supports up to 128 documents per request.
    """
    texts = [open(p).read() for p in paths]
    all_embeddings = []

    for i in range(0, len(texts), batch_size):
        batch  = texts[i : i + batch_size]
        result = client.embed(batch, model=model, input_type="document")
        all_embeddings.extend(result.embeddings)

    return all_embeddings


def embed_query(query: str, model: str = "voyage-4-large") -> list[float]:
    result = client.embed([query], model=model, input_type="query")
    return result.embeddings[0]

from sentence_transformers import SentenceTransformer
import torch

model = SentenceTransformer(
    "nvidia/llama-embed-nemotron-8b",
    device="cuda",
    trust_remote_code=True,
)

# Nemotron uses an instruction prefix for asymmetric retrieval.
# Documents are embedded without a prefix; queries use this one.
QUERY_PREFIX = "Instruct: Retrieve relevant Solidity code\nQuery: "

def embed_documents(texts: list[str], batch_size: int = 16) -> torch.Tensor:
    return model.encode(
        texts,
        batch_size=batch_size,
        convert_to_tensor=True,
        normalize_embeddings=True,
        show_progress_bar=True,
    )

def embed_query_nemotron(query: str) -> torch.Tensor:
    return model.encode(
        QUERY_PREFIX + query,
        convert_to_tensor=True,
        normalize_embeddings=True,
    )

from sentence_transformers import SentenceTransformer
import torch

model = SentenceTransformer(
    "bflhc/Octen-Embedding-8B",
    device="cuda",
    trust_remote_code=True,
)

def embed_batch(texts: list[str], batch_size: int = 8) -> torch.Tensor:
    return model.encode(
        texts,
        batch_size=batch_size,
        convert_to_tensor=True,
        normalize_embeddings=True,
        show_progress_bar=True,
    )

from sentence_transformers import SentenceTransformer
import numpy as np

# Loads entirely into CPU RAM; ~607 MB
model = SentenceTransformer(
    "Shuu12121/CodeSearch-ModernBERT-Crow-Plus",
    device="cpu",
)

def embed(texts: list[str], batch_size: int = 64) -> np.ndarray:
    return model.encode(
        texts,
        batch_size=batch_size,
        normalize_embeddings=True,
        show_progress_bar=False,
    )

import re
from dataclasses import dataclass, field


@dataclass
class SolidityChunk:
    contract:   str
    name:       str
    kind:       str   # "function" | "modifier" | "constructor" | etc.
    source:     str
    start_line: int
    metadata:   dict = field(default_factory=dict)


def chunk_by_function(source: str, contract_name: str = "") -> list[SolidityChunk]:
    """
    Naive regex-based function extractor.
    For production use, prefer a proper Solidity AST parser
    (e.g. solidity-parser-antlr via subprocess, or slither's AST walker).
    """
    pattern = re.compile(
        r"""
        (?P<kind>function|modifier|constructor|receive|fallback)
        \s+
        (?P<name>\w+)?   # constructors / receive / fallback have no name
        [^{]*            # visibility, returns, override, modifier list
        \{
        """,
        re.VERBOSE,
    )

    chunks = []
    for m in pattern.finditer(source):
        start_pos  = m.start()
        start_line = source[:start_pos].count("\n") + 1

        # Walk forward to find the matching closing brace
        depth = 0
        end   = start_pos
        for i, ch in enumerate(source[start_pos:], start_pos):
            if ch == "{":
                depth += 1
            elif ch == "}":
                depth -= 1
                if depth == 0:
                    end = i + 1
                    break

        kind = m.group("kind")
        name = m.group("name") or kind  # unnamed receive/fallback

        chunks.append(SolidityChunk(
            contract   = contract_name,
            name       = name,
            kind       = kind,
            source     = source[start_pos:end].strip(),
            start_line = start_line,
        ))

    return chunks

from transformers import AutoTokenizer
import numpy as np

tokenizer = AutoTokenizer.from_pretrained("Shuu12121/CodeSearch-ModernBERT-Crow-Plus")


def sliding_window_chunks(
    text: str,
    max_tokens: int = 900,  # leave headroom under the 1024 hard limit
    overlap:    int = 100,
) -> list[str]:
    """
    Tokenise -> split into overlapping windows -> decode back to strings.
    Overlap ensures context isn't lost at seam boundaries.
    """
    ids  = tokenizer.encode(text, add_special_tokens=False)
    step = max_tokens - overlap
    windows = []
    for start in range(0, len(ids), step):
        window = ids[start : start + max_tokens]
        windows.append(tokenizer.decode(window))
        if start + max_tokens >= len(ids):
            break
    return windows


def embed_contract_windowed(
    source:   str,
    embed_fn,          # callable: list[str] -> np.ndarray
    pool:     str = "mean",
) -> np.ndarray:
    """
    Embed a full contract via sliding windows, then pool chunk vectors
    into a single representative vector.
    """
    chunks = sliding_window_chunks(source)
    vecs   = embed_fn(chunks)   # shape: (n_chunks, dim)
    if pool == "mean":
        return vecs.mean(axis=0)
    elif pool == "max":
        return vecs.max(axis=0)
    raise ValueError(f"Unknown pool strategy: {pool}")

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer(
    "jinaai/jina-embeddings-v5-text-small",
    device="cpu",
    trust_remote_code=True,
)

def embed_whole_file(sol_path: str) -> np.ndarray:
    """
    With a 32k context window we can embed the entire file in one shot.
    At ~4 chars/token a 120k-character contract still fits comfortably.
    """
    source = open(sol_path).read()
    return model.encode(source, normalize_embeddings=True)

from sentence_transformers import SentenceTransformer

model = SentenceTransformer(
    "bflhc/Octen-Embedding-0.6B",
    device="cpu",
    trust_remote_code=True,
)

from sentence_transformers import SentenceTransformer
import torch

model = SentenceTransformer(
    "bflhc/Octen-Embedding-4B",
    device="cuda",
    trust_remote_code=True,
    model_kwargs={"torch_dtype": torch.float16},  # halves VRAM to ~3.8 GB
)

def embed_batch(texts: list[str], batch_size: int = 32) -> torch.Tensor:
    return model.encode(
        texts,
        batch_size=batch_size,
        convert_to_tensor=True,
        normalize_embeddings=True,
        show_progress_bar=True,
    )

from sentence_transformers import SentenceTransformer
import torch

model = SentenceTransformer(
    "Qwen/Qwen3-Embedding-4B",
    device="cuda",
    trust_remote_code=True,
    model_kwargs={"torch_dtype": torch.bfloat16},
)

# Qwen3 embedding models support an asymmetric task instruction prefix.
# Documents are encoded without any prefix; queries use the one below.
TASK_INSTRUCTION = (
    "Instruct: Given a natural language query about Solidity smart contracts, "
    "retrieve the most relevant code.\nQuery: "
)

def embed_documents_qwen(texts: list[str]) -> torch.Tensor:
    return model.encode(
        texts,
        convert_to_tensor=True,
        normalize_embeddings=True,
        prompt_name=None,
    )

def embed_query_qwen(query: str) -> torch.Tensor:
    return model.encode(
        TASK_INSTRUCTION + query,
        convert_to_tensor=True,
        normalize_embeddings=True,
    )

import numpy as np
import hnswlib
import pickle
from pathlib import Path
from dataclasses import dataclass


@dataclass
class IndexedChunk:
    path:       str
    contract:   str
    name:       str
    kind:       str
    source:     str
    start_line: int


class SolidityIndex:
    """
    Thin wrapper around hnswlib for Solidity code retrieval.
    Agnostic to embedding model — pass in any embed_fn.
    """

    def __init__(self, dim: int, space: str = "cosine", max_elements: int = 100_000):
        self.dim   = dim
        self.index = hnswlib.Index(space=space, dim=dim)
        self.index.init_index(max_elements=max_elements, ef_construction=200, M=16)
        self.index.set_ef(50)
        self.chunks: list[IndexedChunk] = []

    def add(self, chunks: list[IndexedChunk], embeddings: np.ndarray) -> None:
        start = len(self.chunks)
        ids   = np.arange(start, start + len(chunks))
        self.index.add_items(embeddings, ids)
        self.chunks.extend(chunks)

    def search(
        self,
        query_vec: np.ndarray,
        k: int = 5,
    ) -> list[tuple[IndexedChunk, float]]:
        labels, distances = self.index.knn_query(query_vec.reshape(1, -1), k=k)
        return [
            (self.chunks[label], float(dist))
            for label, dist in zip(labels[0], distances[0])
        ]

    def save(self, path: str) -> None:
        p = Path(path)
        self.index.save_index(str(p.with_suffix(".bin")))
        with open(p.with_suffix(".meta"), "wb") as f:
            pickle.dump(self.chunks, f)

    @classmethod
    def load(cls, path: str, dim: int) -> "SolidityIndex":
        p   = Path(path)
        idx = cls(dim=dim)
        idx.index.load_index(str(p.with_suffix(".bin")), max_elements=100_000)
        with open(p.with_suffix(".meta"), "rb") as f:
            idx.chunks = pickle.load(f)
        return idx

from sentence_transformers import SentenceTransformer
from pathlib import Path

embed_model = SentenceTransformer(
    "Shuu12121/CodeSearch-ModernBERT-Crow-Plus", device="cpu"
)

def build_index(sol_dir: str, index_path: str = "solidity.index") -> SolidityIndex:
    sol_files  = list(Path(sol_dir).rglob("*.sol"))
    all_chunks: list[IndexedChunk] = []

    for sol_file in sol_files:
        source = sol_file.read_text(errors="replace")
        for c in chunk_by_function(source, sol_file.stem):
            all_chunks.append(IndexedChunk(
                path=str(sol_file), contract=c.contract, name=c.name,
                kind=c.kind, source=c.source, start_line=c.start_line,
            ))

    print(f"Indexing {len(all_chunks)} chunks from {len(sol_files)} files...")
    vecs = embed_model.encode(
        [c.source for c in all_chunks],
        batch_size=64,
        normalize_embeddings=True,
        show_progress_bar=True,
    )

    idx = SolidityIndex(dim=vecs.shape[1])
    idx.add(all_chunks, vecs)
    idx.save(index_path)
    return idx


def query_index(query: str, index_path: str = "solidity.index", k: int = 5) -> None:
    idx   = SolidityIndex.load(index_path, dim=768)
    q_vec = embed_model.encode(query, normalize_embeddings=True)

    print(f'\nTop {k} for: "{query}"\n{"─" * 60}')
    for rank, (chunk, dist) in enumerate(idx.search(q_vec, k=k), 1):
        print(f"[{rank}] {chunk.contract}::{chunk.name}  ({chunk.kind})")
        print(f"     {chunk.path}:{chunk.start_line}  •  cosine dist: {dist:.4f}")
        print(f"     {chunk.source[:120].strip()}...")
        print()

# Index a contracts directory
python -c "from retrieval import build_index; build_index('./contracts')"

# Query
python -c "
from retrieval import query_index
query_index('reentrancy guard on withdraw function')
query_index('fee-on-transfer token handling')
query_index('emergency pause with time delay')
"

import json
import numpy as np
import pandas as pd
from pathlib import Path
from sentence_transformers import SentenceTransformer


def load_eval_set(path: str) -> list[dict]:
    """
    Expected format — a JSON array:
    [
      {
        "query": "reentrancy guard on ETH withdrawal",
        "relevant": ["Vault::withdraw", "LendingPool::repay"],
        "hard_negatives": ["Vault::deposit"]
      },
      ...
    ]
    Label hard negatives carefully: superficially similar code that is
    functionally different is the failure mode you're testing against.
    """
    return json.loads(Path(path).read_text())


def recall_at_k(
    model:         SentenceTransformer,
    corpus_chunks: list[IndexedChunk],
    corpus_vecs:   np.ndarray,
    eval_set:      list[dict],
    k:             int = 5,
) -> float:
    """
    Compute Recall@K across the eval set.
    A query is a hit if at least one relevant chunk appears in the top-K results.
    Vectors must be L2-normalised (cosine sim == dot product).
    """
    hits = 0
    for item in eval_set:
        q_vec = model.encode(item["query"], normalize_embeddings=True)
        sims  = corpus_vecs @ q_vec
        top_k = np.argsort(sims)[::-1][:k]

        retrieved = {
            f"{corpus_chunks[i].contract}::{corpus_chunks[i].name}"
            for i in top_k
        }
        if any(r in retrieved for r in item["relevant"]):
            hits += 1

    return hits / len(eval_set)


def compare_models(
    model_ids:     list[str],
    corpus_chunks: list[IndexedChunk],
    eval_set:      list[dict],
    ks:            list[int] = [1, 5, 10],
) -> pd.DataFrame:
    """
    Run every candidate model through the eval set and return a comparison
    table, sorted by R@5. Run this before committing to a full re-index.
    """
    texts = [c.source for c in corpus_chunks]
    rows  = []

    for model_id in model_ids:
        print(f"Evaluating {model_id}...")
        m    = SentenceTransformer(model_id, device="cpu", trust_remote_code=True)
        vecs = m.encode(texts, normalize_embeddings=True, show_progress_bar=False)
        row  = {"model": model_id}
        for k in ks:
            row[f"R@{k}"] = recall_at_k(m, corpus_chunks, vecs, eval_set, k=k)
        rows.append(row)

    return pd.DataFrame(rows).sort_values("R@5", ascending=False)

SELECTED = [
    "voyage-4-large (embed_dim=2048)",
    "llama-embed-nemotron-8b",
    "Octen-Embedding-8B",
    "Octen-Embedding-4B",
    "MoD-Embedding",
    "CodeSearch-ModernBERT-Crow-Plus",
    "jina-embeddings-v5-text-small",
    "Octen-Embedding-0.6B",
    "Qwen3-Embedding-4B",
    "C2LLM-7B",
]

display_cols = [
    "Model", "Retrieval",
    "Memory Usage (MB)", "Number of Parameters (B)",
    "Embedding Dimensions", "Max Tokens",
    "CodeSearchNetCCRetrieval", "CodeSearchNetRetrieval",
]

matrix = (
    merged[merged["Model"].isin(SELECTED)][display_cols]
    .sort_values("Retrieval", ascending=False)
    .rename(columns={
        "Memory Usage (MB)":          "Mem (MB)",
        "Number of Parameters (B)":   "Params (B)",
        "Embedding Dimensions":        "Embed Dim",
        "Max Tokens":                  "Max Tok",
        "CodeSearchNetCCRetrieval":    "CSNCC ★",
        "CodeSearchNetRetrieval":      "CSN ★",
    })
)

print(matrix.to_markdown(index=False, floatfmt=".2f"))