File size: 2,730 Bytes
b80417a | 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 | """Predict topics from raw text — no training stack required.
Loads the two artifacts this repo ships and the public base model:
weights/seed0/lora/ LoRA adapter (~10.1M params) -> corrects the encoder
weights/seed0/head/head.pt 56 prototypes + scale -> turns vectors into topics
Qwen/Qwen3-Embedding-0.6B NOT stored here; transformers downloads it
from huggingface.co on first load (that's the point:
we ship only the parts we created).
Usage:
python predict.py --weights weights/seed0 "Cold brew, oat milk, double"
echo "PSA: picture day moved to Friday" | python predict.py --weights weights/seed0
"""
import argparse
import json
import sys
import torch
import torch.nn.functional as F
from peft import PeftModel
from transformers import AutoModel, AutoTokenizer
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
def load(weights_dir):
hs = torch.load(f"{weights_dir}/head/head.pt", map_location="cpu")
tok = AutoTokenizer.from_pretrained(hs["base_model"])
tok.padding_side = "right"
base = AutoModel.from_pretrained(hs["base_model"], torch_dtype=torch.bfloat16).to(DEVICE)
base.config.use_cache = False
model = PeftModel.from_pretrained(base, f"{weights_dir}/lora").to(DEVICE).eval()
protos = F.normalize(hs["protos"].to(DEVICE).float(), dim=-1)
return tok, model, protos, float(hs["scale"]), hs["topics"]
@torch.no_grad()
def encode(model, tok, texts, max_length=512):
enc = tok(texts, padding=True, truncation=True, max_length=max_length,
return_tensors="pt").to(DEVICE)
out = model(**enc).last_hidden_state
last = enc["attention_mask"].sum(1) - 1
return out[torch.arange(out.size(0), device=DEVICE), last].float()
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--weights", default="weights/seed0")
ap.add_argument("--topk", type=int, default=3)
ap.add_argument("text", nargs="?", default=None)
a = ap.parse_args()
texts = [a.text] if a.text else [l.strip() for l in sys.stdin if l.strip()]
tok, model, protos, scale, topics = load(a.weights)
v = F.normalize(encode(model, tok, texts), dim=-1)
probs = F.softmax(scale * (v @ protos.t()), dim=-1)
for t, p in zip(texts, probs):
top = torch.topk(p, a.topk)
print(json.dumps({
"text": (t[:70] + "…") if len(t) > 70 else t,
"topic": topics[top.indices[0]],
"confidence": round(float(top.values[0]), 4),
"top3": [{"topic": topics[i], "p": round(float(x), 4)}
for i, x in zip(top.indices, top.values)],
}, ensure_ascii=False))
if __name__ == "__main__":
main()
|