wehub-resource-sync
a06f331eb8
install-script / powershell-syntax (push) Waiting to run
CI / benchmark (push) Has been skipped
install-script / posix-syntax (push) Successful in 6m1s
install-script / install (macos-14) (push) Waiting to run
install-script / install (ubuntu-latest) (push) Waiting to run
CI / build-onnx (push) Failing after 6m43s
init-smoke / dry-run (push) Failing after 15m57s
security / govulncheck (push) Has been cancelled
security / trivy-fs (push) Has been cancelled
CI / test (1.26, ubuntu-latest) (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
CI / test (1.26, macos-latest) (push) Has been cancelled
CI / build-windows (push) Has been cancelled
CI / lint (push) Has been cancelled
93 lines
2.8 KiB
Python
93 lines
2.8 KiB
Python
#!/usr/bin/env python3
|
|
"""CodeRankEmbed Hybrid baseline runner.
|
|
|
|
Wrapper that lets the Go-side bench/baselines harness invoke
|
|
CodeRankEmbed Hybrid without per-baseline Go code growing model-
|
|
download logic. Usage:
|
|
|
|
python3 bench/baselines/python/coderankembed_runner.py \\
|
|
--repo PATH --query "validateToken" --top-k 10
|
|
|
|
Emits one repo-relative path per line on stdout. Errors go to
|
|
stderr; non-zero exit when the model isn't available.
|
|
|
|
Install: `pip install sentence-transformers transformers torch`.
|
|
First run downloads the CodeRankEmbed model (~440 MB).
|
|
"""
|
|
|
|
import argparse
|
|
import os
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
def main() -> int:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--repo", required=True, help="indexed corpus path")
|
|
ap.add_argument("--query", required=True, help="single query string")
|
|
ap.add_argument("--top-k", type=int, default=10)
|
|
args = ap.parse_args()
|
|
|
|
try:
|
|
from sentence_transformers import SentenceTransformer
|
|
except ImportError as e:
|
|
print(
|
|
f"coderankembed_runner: missing dependency ({e}). "
|
|
"pip install sentence-transformers transformers torch",
|
|
file=sys.stderr,
|
|
)
|
|
return 2
|
|
|
|
model = SentenceTransformer("nomic-ai/CodeRankEmbed", trust_remote_code=True)
|
|
|
|
# Index every file under repo (cheap for sub-million LoC; the
|
|
# ground-truth fixture is the gortex repo itself).
|
|
repo = Path(args.repo).resolve()
|
|
paths: list[Path] = []
|
|
texts: list[str] = []
|
|
for p in repo.rglob("*"):
|
|
if not p.is_file():
|
|
continue
|
|
if any(seg.startswith(".") for seg in p.relative_to(repo).parts):
|
|
continue
|
|
if p.suffix.lower() not in {
|
|
".go", ".py", ".ts", ".tsx", ".js", ".jsx", ".rs",
|
|
".java", ".kt", ".swift", ".rb", ".cs", ".cpp", ".c",
|
|
".h", ".hpp", ".md", ".yaml", ".yml", ".json",
|
|
}:
|
|
continue
|
|
try:
|
|
text = p.read_text(errors="ignore")
|
|
except OSError:
|
|
continue
|
|
if not text.strip():
|
|
continue
|
|
paths.append(p)
|
|
texts.append(text[:8000]) # truncate to keep the embed cheap
|
|
|
|
if not texts:
|
|
print(
|
|
"coderankembed_runner: no indexable files under repo",
|
|
file=sys.stderr,
|
|
)
|
|
return 1
|
|
|
|
embeds = model.encode(texts, show_progress_bar=False, convert_to_numpy=True)
|
|
qe = model.encode([args.query], show_progress_bar=False, convert_to_numpy=True)[0]
|
|
|
|
# Cosine similarity → rank.
|
|
import numpy as np
|
|
|
|
sims = embeds @ qe / (
|
|
(np.linalg.norm(embeds, axis=1) * np.linalg.norm(qe)) + 1e-12
|
|
)
|
|
order = np.argsort(-sims)[: args.top_k]
|
|
for idx in order:
|
|
rel = paths[idx].relative_to(repo)
|
|
print(rel.as_posix())
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|