File size: 4,961 Bytes
79ecbd3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 | """Shared pieces for the eval scripts: model loading, generation, a numpy
retriever, and results I/O. Paths and the generator are configured through
environment variables (see example.env; copy it to .env and edit).
"""
import os, json, time, gc, re
os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
def _load_dotenv(path=".env"):
"""Read KEY=value lines into the environment (existing values win)."""
try:
for line in open(path):
line = line.strip()
if not line or line.startswith("#") or "=" not in line:
continue
k, _, v = line.partition("=")
if v.strip():
os.environ.setdefault(k.strip(), v.strip().strip("'\""))
except FileNotFoundError:
pass
_load_dotenv()
import numpy as np
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
# MODEL_TAG keeps comparison-model output files separate from earlier runs.
GEN_MODEL = os.environ.get("GEN_MODEL", "Qwen/Qwen3-4B-Instruct-2507")
MODEL_TAG = os.environ.get("MODEL_TAG", "")
EMB_MODEL = "BAAI/bge-base-en-v1.5"
BGE_QUERY_PREFIX = "Represent this sentence for searching relevant passages: "
# Same containment preamble as Check-In 3. The revised entailment preamble
# is intentionally NOT used here; see decision log D-16.
SYSTEM_PREAMBLE = (
"You are a regulatory assistant. Answer the question using ONLY the "
"context passages below. Cite the supporting document for every claim. "
"If the context does not contain the answer, reply: \"The provided "
"sources do not answer this question\" and briefly say what source "
"likely would."
)
NORAG_PREAMBLE = (
"You are a regulatory assistant. Answer the question from your own "
"knowledge. If you do not know, say so."
)
_PIPE = None
def get_pipe(model_name=GEN_MODEL):
global _PIPE
if _PIPE is None:
tok = AutoTokenizer.from_pretrained(model_name)
mdl = AutoModelForCausalLM.from_pretrained(
model_name, device_map="auto", dtype=torch.bfloat16)
_PIPE = pipeline("text-generation", model=mdl, tokenizer=tok)
return _PIPE
def generate(user_block, preamble, tokens=400):
pipe = get_pipe()
messages = [{"role": "user", "content": preamble + "\n\n" + user_block}]
out = pipe(messages, max_new_tokens=tokens, do_sample=False,
return_full_text=False)
return out[0]["generated_text"].strip()
class Retriever:
"""Embed-and-search over passage dicts; caches the embedding matrix."""
def __init__(self, passages, cache_path):
from sentence_transformers import SentenceTransformer
self.passages = passages
self.model = SentenceTransformer(
EMB_MODEL, device="cuda" if torch.cuda.is_available() else "cpu")
if os.path.exists(cache_path):
self.mat = np.load(cache_path)
assert self.mat.shape[0] == len(passages), (
f"cache {cache_path} has {self.mat.shape[0]} rows but "
f"{len(passages)} passages were loaded; delete the cache")
else:
self.mat = self.model.encode(
[p["text"] for p in passages], normalize_embeddings=True,
batch_size=64, show_progress_bar=True, convert_to_numpy=True)
os.makedirs(os.path.dirname(cache_path), exist_ok=True)
np.save(cache_path, self.mat)
def search(self, query, k=5):
q = self.model.encode([BGE_QUERY_PREFIX + query],
normalize_embeddings=True,
convert_to_numpy=True)[0]
scores = self.mat @ q
idx = np.argsort(-scores)[:k]
return [(self.passages[i], float(scores[i])) for i in idx]
def context_block(hits):
return "\n\n".join(
f"[{j+1}] ({h['id']}): \"{h['text']}\"" for j, (h, _) in enumerate(hits))
def norm_ws(s):
return re.sub(r"\s+", " ", s).strip().lower()
def save_outputs(out_dir, name, mode, metrics, samples, extra=None):
os.makedirs(out_dir, exist_ok=True)
results = {"benchmark": name, "mode": mode, "model": GEN_MODEL,
"embedding": EMB_MODEL, "n_items": len(samples),
"metrics": metrics, "finished": time.strftime("%Y-%m-%d %H:%M")}
if extra:
results.update(extra)
stem = f"{name}_{mode}" + (f"_{MODEL_TAG}" if MODEL_TAG else "")
rp = os.path.join(out_dir, f"{stem}_results.json")
sp = os.path.join(out_dir, f"{stem}_samples.json")
with open(rp, "w") as fh:
json.dump(results, fh, indent=2)
with open(sp, "w") as fh:
json.dump(samples, fh, indent=2)
print(f"saved {rp}\nsaved {sp}")
print(json.dumps(metrics, indent=2))
def cleanup():
global _PIPE
_PIPE = None
gc.collect()
if torch.cuda.is_available():
torch.cuda.empty_cache()
|