| """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() |
|
|
