TanujInsane/document-classification-env
3
1"""2Generate realistic synthetic customer support and internal escalation documents.3Features multi-format data (JSON, Slack, Logs, Emails), ambiguity, and VIP handling.4"""5 6import numpy as np7import random8import string9import json10from typing import Tuple, List, Dict11from sklearn.feature_extraction.text import TfidfVectorizer12from constants import CATEGORIES13 14 15def get_task_config(difficulty: str) -> Dict:16 configs = {17 "easy": {18 "num_documents": 100,19 "num_categories": 5,20 "feature_dim": 100,21 "max_word_count": 200,22 "time_limit_per_action": None,23 "time_limit_per_episode": None,24 "ambiguity_rate": 0.0,25 "sla_seconds": None,26 },27 "medium": {28 "num_documents": 500,29 "num_categories": 10,30 "feature_dim": 100,31 "max_word_count": 500,32 "time_limit_per_action": 2.0,33 "time_limit_per_episode": 1800,34 "ambiguity_rate": 0.15,35 "sla_seconds": 120.0,36 },37 "hard": {38 "num_documents": 1000,39 "num_categories": 22,40 "feature_dim": 100,41 "max_word_count": 1000,42 "time_limit_per_action": 1.0,43 "time_limit_per_episode": 1200,44 "ambiguity_rate": 0.30,45 "sla_seconds": 60.0,46 }47 }48 return configs[difficulty]49 50FIRST_NAMES = ["James", "Sarah", "Michael", "Emily", "David", "Jessica", "Robert", "Wei", "Fatima", "Carlos"]51LAST_NAMES = ["Smith", "Johnson", "Williams", "Brown", "Jones", "Garcia", "Miller", "Chen", "Al-Rashid"]52COMPANIES = ["Acme Corp", "TechForward Inc", "Global Solutions Ltd", "Stripe", "AWS", "Datadog"]53 54def _ctx():55 name = f"{random.choice(FIRST_NAMES)} {random.choice(LAST_NAMES)}"56 is_vip = random.random() < 0.15 # 15% chance of being a VIP enterprise client57 return {58 "name": name,59 "email": f"{name.lower().replace(' ', '.')}@{random.choice(['gmail.com', 'company.com', 'enterprise.io'])}",60 "company": random.choice(COMPANIES),61 "order_id": f"ORD-{random.randint(100000, 999999)}",62 "amount": f"${random.randint(50, 9999)}.{random.randint(0, 99):02d}",63 "date": f"2025-{random.randint(1,12):02d}-{random.randint(1,28):02d}",64 "is_vip": is_vip,65 "urgency_level": "CRITICAL" if is_vip else random.choice(["LOW", "MEDIUM", "HIGH"])66 }67 68class TaskDataGenerator:69 def __init__(self, difficulty: str, seed: int = None):70 self.difficulty = difficulty71 self.config = get_task_config(difficulty)72 if seed is not None:73 np.random.seed(seed)74 random.seed(seed)75 76 self.categories = list(CATEGORIES[difficulty].values())77 78 # Simple vocabulary for TF-IDF training79 self.vectorizer = TfidfVectorizer(max_features=self.config["feature_dim"], lowercase=True, stop_words='english')80 self.vectorizer.fit(["invoice billing refund", "technical bug crash", "legal contract compliance", "hr payroll benefits complaint", "marketing operations finance strategy tax social"])81 82 def _get_corpus_text(self) -> str:83 if not hasattr(self, "_hf_corpus"):84 try:85 from datasets import load_dataset86 ds = load_dataset("ag_news", split="test")87 self._hf_corpus = [x["text"] for x in ds]88 except Exception:89 self._hf_corpus = ["Please assist with process.", "Error encountered in system."]90 return random.choice(self._hf_corpus)91 92 def _generate_format(self, category: str, ctx: dict) -> str:93 """Generates content in various formats (JSON, Slack, Log, Email) with Real Corpus Injection"""94 fmt = random.choice(["json", "slack", "log", "email"])95 cat_lower = category.lower()96 real_text = self._get_corpus_text()97 98 content = ""99 if fmt == "json":100 payload = {101 "event_id": f"evt_{random.randint(1000,9999)}",102 "source": "webhook_listener",103 "timestamp": ctx["date"],104 "customer": {"id": ctx["email"], "vip_tier": "Enterprise" if ctx["is_vip"] else "Standard"},105 "payload": f"Category: {cat_lower}. Order {ctx['order_id']} value {ctx['amount']}. Extract: {real_text}",106 "metadata": {"raw_category_hint": category, "severity": ctx["urgency_level"]}107 }108 content = json.dumps(payload, indent=2)109 110 elif fmt == "slack":111 content = f"**Slack Transcript - Channel: #triage-incoming**\n"112 content += f"[{ctx['date']} 09:12 AM] @system_bot: New ticket from {ctx['name']} ({ctx['company']}). VIP: {ctx['is_vip']}\n"113 content += f"[{ctx['date']} 09:13 AM] {ctx['name']}: Hey team, we have an issue regarding {cat_lower}.\n"114 content += f"[{ctx['date']} 09:14 AM] {ctx['name']}: For context: {real_text}"115 116 elif fmt == "log":117 content = f"2025-01-01T12:00:00Z WARN [req_id={random.randint(100,999)}] subsystem={category.split('-')[0]} action=PROCESS userId={ctx['email']} msg=\"{real_text}\""118 119 else: # Email120 vip_tag = "[VIP - URGENT] " if ctx['is_vip'] else ""121 content = f"Subject: {vip_tag}Inquiry regarding {cat_lower}\nFrom: {ctx['email']}\n\nTo whom it may concern,\nWe are reaching out about order {ctx['order_id']} value {ctx['amount']}.\n\nAdditional details: {real_text}\n\nRegards,\n{ctx['name']}"122 123 return content124 125 def generate_task_data(self) -> Tuple[List[Dict], np.ndarray]:126 documents = []127 labels = []128 num_docs = self.config["num_documents"]129 ambiguity_rate = self.config.get("ambiguity_rate", 0.0)130 131 for i in range(num_docs):132 ctx = _ctx()133 134 if random.random() < ambiguity_rate:135 # Deliberate heavy ambiguity: Mention multiple departments136 cat1 = random.choice(self.categories)137 cat2 = random.choice(self.categories)138 true_cat = random.choice([cat1, cat2])139 content = self._generate_format(f"{cat1} and {cat2}", ctx)140 else:141 true_cat = random.choice(self.categories)142 content = self._generate_format(true_cat, ctx)143 144 features = self.vectorizer.transform([content]).toarray()[0]145 if np.max(features) > 0: features = 2 * features / (np.max(features) + 1e-8) - 1146 147 doc = {148 "id": f"doc_{i:06d}",149 "content": content,150 "word_count": len(content.split()),151 "has_urgency_markers": ctx["is_vip"], # VIP flag passed via urgency152 "features": features.tolist(),153 "true_category": true_cat,154 "is_vip": ctx["is_vip"]155 }156 157 documents.append(doc)158 labels.append(self.categories.index(true_cat))159 160 return documents, np.array(labels)161