Upload golden_builder.py

golden_builder.py

@@ -0,0 +1,186 @@
+# golden_builder.py
+# -*- coding: utf-8 -*-
+import json, re, logging, hashlib
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Dict, List, Optional
+from collections import Counter
+
+import numpy as np
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
+log = logging.getLogger("golden-builder")
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+# ========= Utilities =========
+PERSIAN_MAP = {'ك':'ک','ى':'ی','ﻲ':'ی','ﯽ':'ی','أ':'ا','إ':'ا'}
+NOISE = [r"http[s]?://\S+", r"www\.\S+", r"\d{10,}", r"(.)\1{4,}", r"[^\u0600-\u06FF\s\d\.,;:!?()\"'\-]+"]
+
+def clean_text(s: str) -> str:
+    if not isinstance(s, str): return ""
+    for a,b in PERSIAN_MAP.items(): s = s.replace(a,b)
+    for pat in NOISE: s = re.sub(pat, " ", s)
+    s = re.sub(r"\s+", " ", s)
+    s = re.sub(r"\.{2,}", "...", s)
+    s = re.sub(r"\s+([،.;:!?])", r"\1", s)
+    s = re.sub(r"([،.;:!?])(?=[^\s])", r"\1 ", s)
+    return s.strip()
+
+def md5(s: str) -> str:
+    import hashlib as _h
+    return _h.md5(s.encode("utf-8")).hexdigest()
+
+def lex_diversity(s: str) -> float:
+    toks = s.split()
+    return 0.0 if not toks else len(set(toks))/len(toks)
+
+def has_repetition(s: str, n:int=3) -> bool:
+    toks = s.split()
+    if len(toks) < n: return False
+    grams = [tuple(toks[i:i+n]) for i in range(len(toks)-n+1)]
+    from collections import Counter
+    return any(c>2 for c in Counter(grams).values())
+
+# ========= Lightweight NER (regex spans برای متادیتا) =========
+@dataclass
+class LegalEntity:
+    text: str; category: str; start: int; end: int; weight: float
+
+class LegalEntityExtractor:
+    def __init__(self):
+        self._defs = {
+            "STATUTE":  ( [r"قانون\s+(?:اساسی|مدنی|کیفری|کار|تجارت|مجازات)",
+                           r"آیین\s+دادرسی\s+(?:مدنی|کیفری)",
+                           r"ماده\s+\d+", r"تبصره\s+\d+"], 1.0 ),
+            "COURT":    ( [r"دیوان\s+(?:عالی|عدالت)", r"دادگاه\s+(?:عمومی|تجدیدنظر|انقلاب)", r"شعبه\s+\d+"], 0.9 ),
+            "CRIME":    ( [r"کلاهبرداری|اختلاس|ارتشا|خیانت\s+در\s+امانت|جعل|سرقت|قتل"], 1.2 ),
+            "PENALTY":  ( [r"حبس|جزای\s+نقدی|شلاق|قصاص|دیه|محرومیت\s+از\s+حقوق\s+اجتماعی"], 1.1 ),
+            "CIVIL":    ( [r"قرارداد|عقد\s+(?:بیع|اجاره|رهن|نکاح)|خسارت|تعهد|ضمان|مطالبه"], 0.8 ),
+            "PROCED":   ( [r"دادخواست|لایحه|شکواییه|ابلاغ|جلسه\s+دادرسی|کارشناسی|دلایل\s+اثباتی"], 0.7 ),
+            "PARTY":    ( [r"خواهان|خوانده|شاکی|متهم|وکیل\s+دادگستری|دادستان|قاضی"], 0.6 ),
+            "BUSINESS": ( [r"شرکت\s+(?:سهامی|مسئولیت\s+محدود)|ورشکستگی|سهام|چک|سفته|برات"], 0.6 ),
+        }
+        self._patterns = []
+        for cat,(pats,w) in self._defs.items():
+            for p in pats:
+                self._patterns.append( (re.compile(p, re.IGNORECASE), cat, w) )
+
+    def extract(self, text:str) -> List[LegalEntity]:
+        out, seen = [], set()
+        for rgx, cat, w in self._patterns:
+            for m in rgx.finditer(text):
+                s,e = m.span()
+                if (s,e) in seen: continue
+                seen.add((s,e))
+                out.append(LegalEntity(m.group(), cat, s, e, w))
+        out.sort(key=lambda x: x.start)
+        return out
+
+# ========= Builder =========
+class GoldenBuilder:
+    def __init__(self, model_name: str = "google/mt5-base", device: Optional[str] = None,
+                 min_len:int=40, max_len:int=160, min_entities:int=2):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        log.info("Device: %s", self.device)
+        self.tok = AutoTokenizer.from_pretrained(model_name, use_fast=True)
+        self.model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(self.device)
+        self.model.eval()
+        self.min_len, self.max_len = min_len, max_len
+        self.min_entities = min_entities
+        self._seen_hashes = set()
+        self.ner = LegalEntityExtractor()
+
+    def _summarize_batch(self, texts: List[str], num_beams:int=6) -> List[str]:
+        if not texts: return []
+        inputs = self.tok(texts, return_tensors="pt", truncation=True, padding=True, max_length=512).to(self.device)
+        with torch.no_grad():
+            ids = self.model.generate(
+                **inputs,
+                max_length=self.max_len, min_length=self.min_len,
+                num_beams=num_beams, early_stopping=True,
+                length_penalty=2.5, no_repeat_ngram_size=3, do_sample=False
+            )
+        return self.tok.batch_decode(ids, skip_special_tokens=True)
+
+    def _quality_gate(self, src:str, tgt:str, ents:List[LegalEntity]) -> bool:
+        s_len, t_len = len(src.split()), len(tgt.split())
+        if not (30 <= s_len and 20 <= t_len <= 220): return False
+        comp = (t_len/(s_len+1e-8))
+        if not (0.12 <= comp <= 0.65): return False
+        if lex_diversity(tgt) < 0.4: return False
+        if has_repetition(tgt, 3): return False
+        if len(ents) < self.min_entities: return False
+        ent_density = (sum((e.end - e.start) for e in ents) / max(len(src),1)) * 100
+        if ent_density < 4.0: return False
+        return True
+
+    def build(self, raw_items: List[Dict], text_key:str="متن_کامل", batch_size:int=4) -> List[Dict]:
+        rows, i = [], 0
+        N = len(raw_items)
+        while i < N:
+            chunk = raw_items[i:i+batch_size]
+            cleaned = [clean_text(str(it.get(text_key, ""))) for it in chunk]
+            # de-dup + کوتاه‌زدایی
+            todo = []
+            for c in cleaned:
+                if len(c.split()) < 30:
+                    todo.append("")
+                    continue
+                h = md5(c)
+                if h in self._seen_hashes:
+                    todo.append("")
+                    continue
+                self._seen_hashes.add(h)
+                todo.append(c)
+            # Summarize only valid items
+            inputs = [f"summarize: {t}" for t in todo if t]
+            outputs = self._summarize_batch(inputs) if inputs else []
+            k = 0
+            for c in todo:
+                if not c: continue
+                tgt = clean_text(outputs[k]); k += 1
+                ents = self.ner.extract(c)
+                if not self._quality_gate(c, tgt, ents):
+                    continue
+                ents_payload = [{"text": e.text, "category": e.category, "start": e.start, "end": e.end, "weight": e.weight}
+                                for e in ents[:20]]
+                rows.append({
+                    "input": f"summarize: {c}",
+                    "output": tgt,
+                    "metadata": {
+                        "input_length": len(c.split()),
+                        "target_length": len(tgt.split())
+                    },
+                    "legal_entities": {
+                        "total_entities": len(ents),
+                        "categories": dict(Counter(e.category for e in ents)),
+                        "entities": ents_payload
+                    }
+                })
+            i += batch_size
+        return rows
+
+# ========= I/O =========
+def load_json_or_jsonl(path: str) -> List[Dict]:
+    p = Path(path)
+    raw = p.read_text(encoding="utf-8").strip()
+    try:
+        data = json.loads(raw)
+        return data if isinstance(data, list) else [data]
+    except json.JSONDecodeError:
+        out = []
+        for ln in raw.splitlines():
+            ln = ln.strip()
+            if not ln: continue
+            try:
+                out.append(json.loads(ln))
+            except json.JSONDecodeError:
+                pass
+        return out
+
+def save_jsonl(rows: List[Dict], out_path: str):
+    p = Path(out_path); p.parent.mkdir(parents=True, exist_ok=True)
+    with p.open("w", encoding="utf-8") as f:
+        for r in rows:
+            f.write(json.dumps(r, ensure_ascii=False) + "\n")