| |
| """Tag one Saudi government tender: Arabic title (+ agency) in, labels and confidences out. |
| |
| python predict.py --title "توريد أجهزة حاسب آلي محمول" --agency "وزارة الصحة" |
| python predict.py --title "..." --agency "..." --json |
| python predict.py --file tenders.csv # columns: title, agency -> prints JSON lines |
| |
| Outputs: Etimad activity category (top 3), procurement type, sector (derived from the |
| predicted activity), kind supply/works/services (rules on activity, type and title). |
| Models are chosen per label in models/saudi-tender-tagger/config.json.""" |
| import argparse, csv, json, os, sys, time |
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, os.path.join(HERE, "scripts")) |
| from common import make_input, sector_of, kind_of |
|
|
| MODEL_DIR = os.path.join(HERE, "models", "saudi-tender-tagger") |
|
|
| class Head: |
| def __init__(self, spec): |
| self.kind, path = spec["type"], os.path.join(MODEL_DIR, spec["path"]) |
| if self.kind == "encoder": |
| import torch |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification |
| self.torch = torch |
| self.dev = "mps" if torch.backends.mps.is_available() else "cpu" |
| self.tok = AutoTokenizer.from_pretrained(path) |
| self.model = AutoModelForSequenceClassification.from_pretrained(path).to(self.dev).eval() |
| self.labels = [self.model.config.id2label[i] for i in range(len(self.model.config.id2label))] |
| else: |
| import joblib |
| self.wv, self.cv, self.clf = joblib.load(os.path.join(path, "model.joblib")) |
| self.labels = list(self.clf.classes_) |
| def probs(self, text): |
| if self.kind == "encoder": |
| b = self.tok([text], truncation=True, max_length=128, return_tensors="pt").to(self.dev) |
| with self.torch.no_grad(): |
| return self.torch.softmax(self.model(**b).logits[0], -1).cpu().tolist() |
| from scipy.sparse import hstack |
| return self.clf.predict_proba(hstack([self.wv.transform([text]), self.cv.transform([text])]).tocsr())[0].tolist() |
| def top(self, text, k=1): |
| p = self.probs(text) |
| return sorted(zip(self.labels, p), key=lambda x: -x[1])[:k] |
|
|
| def load_heads(): |
| cfg = json.load(open(os.path.join(MODEL_DIR, "config.json"))) |
| return {k: Head(v) for k, v in cfg["heads"].items()} |
|
|
| def tag(heads, title, agency): |
| text = make_input(title, agency) |
| t0 = time.perf_counter() |
| act = heads["activity"].top(text, 3); tt = heads["ttype"].top(text, 1)[0] |
| activity = act[0][0] |
| if "sector" in heads: |
| sec = heads["sector"].top(text, 1)[0]; sector = {"label": sec[0], "confidence": round(sec[1], 4), "derived_from": "direct sector classifier"} |
| else: |
| sector = {"label": sector_of(activity), "derived_from": "predicted activity"} |
| out = { |
| "activity": {"label": activity, "confidence": round(act[0][1], 4), |
| "top3": [{"label": l, "confidence": round(p, 4)} for l, p in act]}, |
| "procurement_type": {"label": tt[0], "confidence": round(tt[1], 4)}, |
| "sector": sector, |
| "kind": {"label": kind_of(activity, tt[0], title), "derived_from": "rules on activity, type and title"}, |
| } |
| out["latency_ms"] = round((time.perf_counter() - t0) * 1000, 1) |
| return out |
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--title"); ap.add_argument("--agency", default=""); ap.add_argument("--file"); ap.add_argument("--json", action="store_true") |
| a = ap.parse_args() |
| if not a.title and not a.file: ap.error("give --title or --file") |
| t0 = time.perf_counter(); heads = load_heads(); load_s = time.perf_counter() - t0 |
| rows = list(csv.DictReader(open(a.file, encoding="utf-8"))) if a.file else [{"title": a.title, "agency": a.agency}] |
| for r in rows: |
| res = tag(heads, r["title"], r.get("agency", "")) |
| if a.json or a.file: print(json.dumps({"title": r["title"], "agency": r.get("agency", ""), **res}, ensure_ascii=False)) |
| else: |
| print(f"Title: {r['title']}\nAgency: {r.get('agency','')}") |
| print("Top-3 activities:") |
| for i, x in enumerate(res["activity"]["top3"], 1): print(f" {i}. {x['label']} p={x['confidence']:.3f}") |
| print(f"Type: {res['procurement_type']['label']} ({res['procurement_type']['confidence']:.2f})") |
| sc = res["sector"].get("confidence"); print(f"Sector: {res['sector']['label']}" + (f" p={sc:.3f}" if sc else "") + f" Kind: {res['kind']['label']}") |
| print(f"Latency: {res['latency_ms']} ms per tender (model load {load_s:.1f} s)") |
| if __name__ == "__main__": main() |
|
|
