129 lines
4.7 KiB
Python
129 lines
4.7 KiB
Python
# =========================================================
|
|
# extract_bert_embeddings.py
|
|
# Pre-compute mean-pooled BERT sentence embeddings for AG News
|
|
# and save them as a .pt tensor.
|
|
#
|
|
# The output file is consumed by drl_mask_policy.py via
|
|
# --embed-file path/to/bert_agnews_8000.pt
|
|
#
|
|
# The consumer script applies dataset-level mean centering and
|
|
# L2 normalization, so the saved tensor here contains the RAW
|
|
# mean-pooled embeddings (no centering, no normalization).
|
|
#
|
|
# Usage:
|
|
# python extract_bert_embeddings.py \
|
|
# --out bert_agnews_8000.pt \
|
|
# --max-sentences 8000
|
|
# =========================================================
|
|
|
|
import argparse
|
|
import os
|
|
import random
|
|
|
|
import torch
|
|
from transformers import BertModel, BertTokenizer
|
|
|
|
|
|
def load_agnews_sentences(max_sentences, min_len=5, max_len=30):
|
|
"""Load first-sentence headlines from AG News. Falls back to
|
|
synthetic templated text if the dataset cannot be fetched."""
|
|
sentences = []
|
|
try:
|
|
from datasets import load_dataset
|
|
ds = load_dataset("ag_news", split="train")
|
|
for example in ds:
|
|
first = example["text"].split(".")[0].strip()
|
|
words = first.split()
|
|
if min_len <= len(words) <= max_len:
|
|
sentences.append(first)
|
|
if len(sentences) >= max_sentences:
|
|
break
|
|
except Exception as e:
|
|
print(f"[WARN] AG News unavailable ({e}). Using synthetic.")
|
|
|
|
if len(sentences) < max_sentences:
|
|
print(f"[INFO] Padding with synthetic sentences "
|
|
f"(have {len(sentences)}, need {max_sentences}).")
|
|
templates = [
|
|
"The {} {} the {} in the {}.",
|
|
"A {} {} quickly {} the {}.",
|
|
"Several {} {} near the {} {}.",
|
|
]
|
|
words = ["system", "signal", "network", "channel", "user",
|
|
"device", "antenna", "receiver", "transmitter",
|
|
"processes", "transmits", "receives", "encodes",
|
|
"wireless", "digital", "robust", "adaptive"]
|
|
while len(sentences) < max_sentences:
|
|
t = random.choice(templates)
|
|
n = t.count("{}")
|
|
sentences.append(t.format(*random.choices(words, k=n)))
|
|
|
|
random.shuffle(sentences)
|
|
return sentences[:max_sentences]
|
|
|
|
|
|
@torch.no_grad()
|
|
def encode_batch(model, tokenizer, texts, device, max_length=64):
|
|
"""Mean-pool token embeddings over non-padding positions."""
|
|
inputs = tokenizer(texts, padding=True, truncation=True,
|
|
max_length=max_length,
|
|
return_tensors="pt").to(device)
|
|
out = model(**inputs)
|
|
hidden = out.last_hidden_state # (B, T, d)
|
|
mask = inputs["attention_mask"].unsqueeze(-1).float()
|
|
summed = (hidden * mask).sum(dim=1) # (B, d)
|
|
count = mask.sum(dim=1).clamp(min=1.0) # (B, 1)
|
|
return (summed / count).cpu() # (B, d)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--model", type=str, default="bert-base-uncased")
|
|
parser.add_argument("--out", type=str, default="bert_agnews_8000.pt")
|
|
parser.add_argument("--max-sentences", type=int, default=8000)
|
|
parser.add_argument("--batch-size", type=int, default=64)
|
|
parser.add_argument("--cuda", action="store_true")
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
args = parser.parse_args()
|
|
|
|
random.seed(args.seed)
|
|
torch.manual_seed(args.seed)
|
|
|
|
if args.cuda and torch.cuda.is_available():
|
|
device = torch.device("cuda")
|
|
elif torch.backends.mps.is_available():
|
|
device = torch.device("mps")
|
|
else:
|
|
device = torch.device("cpu")
|
|
print(f"[INFO] Device: {device}")
|
|
|
|
print(f"[INFO] Loading {args.model} ...")
|
|
tokenizer = BertTokenizer.from_pretrained(args.model)
|
|
model = BertModel.from_pretrained(args.model).to(device)
|
|
model.eval()
|
|
d_bert = model.config.hidden_size
|
|
print(f"[INFO] BERT hidden dim = {d_bert}")
|
|
|
|
print(f"[INFO] Loading {args.max_sentences} AG News sentences ...")
|
|
sents = load_agnews_sentences(args.max_sentences)
|
|
print(f"[INFO] Got {len(sents)} sentences.")
|
|
|
|
all_emb = []
|
|
for i in range(0, len(sents), args.batch_size):
|
|
batch = sents[i:i + args.batch_size]
|
|
emb = encode_batch(model, tokenizer, batch, device)
|
|
all_emb.append(emb)
|
|
if (i // args.batch_size) % 20 == 0:
|
|
print(f"[INFO] Processed {i + len(batch)}/{len(sents)}")
|
|
emb = torch.cat(all_emb, dim=0)
|
|
print(f"[INFO] Final tensor shape: {tuple(emb.shape)}")
|
|
|
|
os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
|
|
torch.save(emb, args.out)
|
|
print(f"[OK] Saved raw mean-pooled embeddings to {args.out}")
|
|
print(f" Feed into drl_mask_policy.py via --embed-file {args.out}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|