Spaces:
Sleeping
Sleeping
| import os, re, json, time, hashlib | |
| from io import BytesIO | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| import pandas as pd | |
| import requests | |
| import streamlit as st | |
| # ───────────────────────────────────────────── | |
| # 1. CONFIGURATION | |
| # ───────────────────────────────────────────── | |
| CATEGORIES = { | |
| "ADMINISTRATIVE": {"desc": "Notifications, circulars, office orders, memos, meeting minutes, policy instructions", | |
| "keywords": ["memo", "circular", "notification", "order", "policy", "minutes", "meeting", "instruction"]}, | |
| "FINANCIAL": {"desc": "Budget, expenditure, audit, payments, bills, revenue, re-appropriation, PC-1 costing", | |
| "keywords": ["budget", "expenditure", "audit", "payment", "revenue", "bill", "fund", "grant", "rupees", "allocation"]}, | |
| "LEGAL": {"desc": "Contracts, agreements, court cases, writ petitions, judgments, legal opinions", | |
| "keywords": ["contract", "agreement", "legal", "court", "judgment", "petition", "writ", "tribunal", "litigation"]}, | |
| "HR": {"desc": "Recruitment, promotion, transfer/posting, leave, pension, disciplinary matters", | |
| "keywords": ["recruitment", "promotion", "transfer", "posting", "leave", "employee", "pension", "appointment", "seniority"]}, | |
| "TECHNICAL": {"desc": "Tenders, specifications, development projects, IT systems, infrastructure", | |
| "keywords": ["specification", "tender", "project", "development", "infrastructure", "procurement", "software", "construction"]}, | |
| "CITIZEN_SERVICE": {"desc": "Citizen applications, complaints, NOCs, requests, grievances", | |
| "keywords": ["complaint", "grievance", "noc", "objection", "applicant", "citizen", "request", "redressal"]}, | |
| } | |
| OPENROUTER_URL = "https://openrouter.ai/api/v1/chat/completions" | |
| DEFAULT_MODEL = "openai/gpt-4o-mini" # check current IDs at openrouter.ai/models | |
| # ───────────────────────────────────────────── | |
| # 2. RULE-BASED ENGINE (optimized v1) | |
| # ───────────────────────────────────────────── | |
| class RuleClassifier: | |
| """Precompiled, word-boundary regex -> no false hits ('order' in 'border').""" | |
| def __init__(self, categories): | |
| self.patterns = { | |
| cat: re.compile(r"\b(?:" + "|".join(map(re.escape, v["keywords"])) + r")\w*\b", re.IGNORECASE) | |
| for cat, v in categories.items() | |
| } | |
| def classify(self, text): | |
| hits = {cat: len(p.findall(text)) for cat, p in self.patterns.items()} | |
| total = sum(hits.values()) | |
| if total == 0: | |
| return {"category": "UNCLASSIFIED", "confidence": 0.0, "scores": hits} | |
| best = max(hits, key=hits.get) | |
| return {"category": best, "confidence": round(hits[best] / total, 2), "scores": hits} | |
| MONTHS = r"(?:Jan|Feb|Mar|Apr|May|Jun|Jul|Aug|Sep|Oct|Nov|Dec)[a-z]*" | |
| DATE_RE = re.compile( | |
| rf"\b(?P<d>\d{{1,2}}(?P<s>[./-])\d{{1,2}}(?P=s)\d{{2,4}}" | |
| rf"|\d{{1,2}}(?:st|nd|rd|th)?\s+{MONTHS},?\s+\d{{4}}" | |
| rf"|{MONTHS}\s+\d{{1,2}},?\s+\d{{4}})\b", re.IGNORECASE) | |
| # Reference must contain a digit -> "No objection" is NOT a reference | |
| REF_RE = re.compile(r"\b(?:Ref(?:erence)?|File|Letter|Endst|No)\.?\s*(?:No\.?)?\s*[:#]?\s*" | |
| r"([A-Z0-9][\w/&().\-]*\d[\w/&().\-]*)", re.IGNORECASE) | |
| def extract_metadata(text): | |
| d, r = DATE_RE.search(text), REF_RE.search(text) | |
| return {"date": d.group("d") if d else None, "reference_no": r.group(1).rstrip(".") if r else None} | |
| # ───────────────────────────────────────────── | |
| # 3. PRIVACY: redact before sending to external AI | |
| # ───────────────────────────────────────────── | |
| PII_PATTERNS = { | |
| "CNIC": re.compile(r"\b\d{5}-?\d{7}-?\d\b"), | |
| "PHONE": re.compile(r"(?:\+92|\b0)3\d{2}[-\s]?\d{7}\b"), | |
| "EMAIL": re.compile(r"[\w.+-]+@[\w-]+\.[\w.]+"), | |
| } | |
| def redact(text): | |
| for label, pattern in PII_PATTERNS.items(): | |
| text = pattern.sub(f"[{label}]", text) | |
| return text | |
| # ───────────────────────────────────────────── | |
| # 4. TEXT EXTRACTION (txt / pdf / docx) | |
| # ───────────────────────────────────────────── | |
| def extract_text(name, data): | |
| ext = name.lower().rsplit(".", 1)[-1] | |
| if ext == "txt": | |
| return data.decode("utf-8", errors="ignore") | |
| if ext == "pdf": | |
| from pypdf import PdfReader | |
| return "\n".join((p.extract_text() or "") for p in PdfReader(BytesIO(data)).pages) | |
| if ext == "docx": | |
| from docx import Document | |
| return "\n".join(p.text for p in Document(BytesIO(data)).paragraphs) | |
| return "" | |
| # ───────────────────────────────────────────── | |
| # 5. AI ENGINE (OpenRouter) | |
| # ───────────────────────────────────────────── | |
| CAT_LIST = "\n".join(f"- {c}: {v['desc']}" for c, v in CATEGORIES.items()) | |
| SYSTEM_PROMPT = f"""You are a records officer AI for the Government of Khyber Pakhtunkhwa, Pakistan. | |
| Classify the document into exactly ONE category: | |
| {CAT_LIST} | |
| Documents may be English, Urdu or mixed, and may contain OCR errors. | |
| Return ONLY valid JSON with these keys: | |
| {{"category": "<one of the categories>", "confidence": <0.0-1.0>, "reason": "<one sentence>", | |
| "subject": "<short subject line>", "summary": "<max 2 sentences>", | |
| "date": "<date or null>", "reference_no": "<file/letter no. or null>", | |
| "department": "<issuing department or null>", "urgency": "Immediate|High|Normal", | |
| "action_required": "<what the receiving officer must do, or null>", "language": "English|Urdu|Mixed"}}""" | |
| def parse_json(content): | |
| content = re.sub(r"```(?:json)?", "", content).strip() | |
| start, end = content.find("{"), content.rfind("}") | |
| return json.loads(content[start:end + 1]) | |
| def llm_classify(text, api_key, model, max_chars=6000): | |
| headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json", | |
| "HTTP-Referer": "http://localhost:8501", "X-Title": "KP Document Classifier"} | |
| payload = { | |
| "model": model, "temperature": 0, | |
| "response_format": {"type": "json_object"}, | |
| "messages": [{"role": "system", "content": SYSTEM_PROMPT}, | |
| {"role": "user", "content": text[:max_chars]}], # truncate = faster + cheaper | |
| } | |
| for attempt in range(3): # retry on rate-limit / server error | |
| r = requests.post(OPENROUTER_URL, headers=headers, json=payload, timeout=60) | |
| if r.status_code == 429 or r.status_code >= 500: | |
| time.sleep(2 ** attempt) | |
| continue | |
| r.raise_for_status() | |
| data = parse_json(r.json()["choices"][0]["message"]["content"]) | |
| if data.get("category") not in CATEGORIES: | |
| data["category"] = "UNCLASSIFIED" | |
| data["confidence"] = max(0.0, min(1.0, float(data.get("confidence", 0)))) | |
| return data | |
| raise RuntimeError("OpenRouter busy — retried 3 times") | |
| # ───────────────────────────────────────────── | |
| # 6. HYBRID PIPELINE | |
| # ───────────────────────────────────────────── | |
| def get_engine(): | |
| return RuleClassifier(CATEGORIES), {} # {} = AI response cache (survives reruns) | |
| def classify_one(name, text, cfg, rules, cache): | |
| t0 = time.perf_counter() | |
| rule, meta = rules.classify(text), extract_metadata(text) | |
| res = {"file": name, "words": len(text.split()), "method": "Rule-based", | |
| "category": rule["category"], "confidence": rule["confidence"], | |
| "rule_scores": rule["scores"], **meta} | |
| if not text.strip(): | |
| res["method"] = "Skipped (no text — run OCR Processor first)" | |
| elif cfg["use_ai"] and cfg["api_key"] and (cfg["always_ai"] or rule["confidence"] < cfg["threshold"]): | |
| safe_text = redact(text) if cfg["redact"] else text | |
| key = hashlib.sha256((cfg["model"] + safe_text).encode()).hexdigest() | |
| try: | |
| ai = cache.get(key) or llm_classify(safe_text, cfg["api_key"], cfg["model"]) | |
| cache[key] = ai | |
| res.update({k: v for k, v in ai.items() if v not in (None, "", "null")}) | |
| res["method"] = "AI (OpenRouter)" | |
| except Exception as e: | |
| res["method"], res["error"] = "Rule-based (AI failed)", str(e) | |
| res["latency_s"] = round(time.perf_counter() - t0, 2) | |
| return res | |
| # ───────────────────────────────────────────── | |
| # 7. STREAMLIT UI | |
| # ───────────────────────────────────────────── | |
| def default_key(): | |
| try: | |
| return st.secrets["OPENROUTER_API_KEY"] | |
| except Exception: | |
| return os.getenv("OPENROUTER_API_KEY", "") | |
| st.set_page_config(page_title="AI Document Classifier v2", page_icon="📄", layout="wide") | |
| st.title("📄 AI Document Classifier v2") | |
| st.caption("Hybrid Rule-based + AI (OpenRouter) • KP Digital Roadmap 2030 — Paperless Government Processes") | |
| with st.sidebar: | |
| st.header("⚙️ Settings") | |
| api_key = default_key() # loaded silently — never rendered | |
| if api_key: | |
| st.success("🔑 API key loaded securely") | |
| else: | |
| api_key = st.text_input("OpenRouter API Key", type="password", | |
| placeholder="Paste key (not saved)") | |
| cfg = { | |
| "api_key": api_key, | |
| "model": st.text_input("Model", value=DEFAULT_MODEL), | |
| "use_ai": st.toggle("Use AI", value=True), | |
| "always_ai": st.toggle("Always use AI (else only when rules unsure)", value=False), | |
| "threshold": st.slider("Rule confidence needed to skip AI", 0.3, 1.0, 0.7, 0.05), | |
| "redact": st.toggle("🔒 Redact CNIC / phone / email before AI", value=True), | |
| "workers": st.slider("Parallel workers", 1, 10, 4), | |
| } | |
| if cfg["use_ai"] and not cfg["api_key"]: | |
| st.warning("No API key — running rule-based only.") | |
| files = st.file_uploader("Upload documents", type=["txt", "pdf", "docx"], accept_multiple_files=True) | |
| if files and st.button("🚀 Classify", type="primary"): | |
| rules, cache = get_engine() | |
| docs = {f.name: extract_text(f.name, f.getvalue()) for f in files} | |
| results, bar, t0 = [], st.progress(0.0, "Classifying..."), time.perf_counter() | |
| with ThreadPoolExecutor(max_workers=cfg["workers"]) as pool: | |
| futures = [pool.submit(classify_one, n, t, cfg, rules, cache) for n, t in docs.items()] | |
| for i, fut in enumerate(as_completed(futures), 1): | |
| results.append(fut.result()) | |
| bar.progress(i / len(futures), f"{i}/{len(futures)} done") | |
| total_time = time.perf_counter() - t0 | |
| df = pd.DataFrame(results) | |
| c1, c2, c3, c4 = st.columns(4) | |
| c1.metric("Documents", len(df)) | |
| c2.metric("AI calls", int(df["method"].eq("AI (OpenRouter)").sum())) | |
| c3.metric("Total time", f"{total_time:.1f}s") | |
| c4.metric("Avg confidence", f"{df['confidence'].mean():.0%}") | |
| cols = [c for c in ["file", "category", "confidence", "method", "subject", "urgency", | |
| "date", "reference_no", "department", "action_required", "latency_s"] if c in df] | |
| st.dataframe(df[cols], use_container_width=True, hide_index=True) | |
| st.bar_chart(df["category"].value_counts()) | |
| low = df[df["confidence"] < 0.6] | |
| if not low.empty: | |
| st.warning(f"👤 {len(low)} document(s) below 60% confidence — route for human review.") | |
| for r in results: | |
| with st.expander(f"📁 {r['file']} → {r['category']} ({r['confidence']:.0%})"): | |
| if r.get("summary"): st.info(r["summary"]) | |
| if r.get("reason"): st.caption(f"Why: {r['reason']}") | |
| if r.get("error"): st.error(r["error"]) | |
| st.json({k: v for k, v in r.items() if k not in ("summary", "reason")}) | |
| st.text_area("Text", docs[r["file"]][:3000], height=150, key=f"t_{r['file']}") | |
| st.download_button("⬇️ Download results (CSV)", df[cols].to_csv(index=False).encode("utf-8"), | |
| "classification_results.csv", "text/csv") |