muffin2006/document-classification-env

1"""2Generate fake documents for training.3Templates for each category, then we mess with them a bit to make variations.4"""5 6import numpy as np7import random8from typing import Tuple, List, Dict9from sklearn.feature_extraction.text import TfidfVectorizer10 11 12def get_task_config(difficulty: str) -> Dict:13    """Just return the config for a difficulty level"""14    configs = {15        "easy": {16            "num_documents": 100,17            "num_categories": 5,18            "feature_dim": 100,19            "max_word_count": 200,20            "time_limit": None,  # Can take as long as you want21        },22        "medium": {23            "num_documents": 500,24            "num_categories": 10,25            "feature_dim": 100,26            "max_word_count": 500,27            "time_limit": 2.0,  # 2 seconds per doc28        },29        "hard": {30            "num_documents": 1000,31            "num_categories": 20,32            "feature_dim": 100,33            "max_word_count": 1000,34            "time_limit": 1.0,  # 1 second, get moving35        }36    }37    return configs[difficulty]38 39 40class TaskDataGenerator:41    """Generate synthetic documents for the task"""42    43    DOCUMENT_TEMPLATES = {44        # Basic category - just regular stuff45        "General": [46            "This is a general inquiry about our services.",47            "I would like to know more about your company.",48            "Could you provide general information?",49            "General question regarding policies.",50        ],51        # Billing stuff52        "Billing": [53            "My invoice shows an incorrect amount. Please review.",54            "I need to update my billing address.",55            "What is the status of my payment?",56            "I have questions about my bill.",57        ],58        # People complaining about being overcharged59        "Billing-Dispute": [60            "I was overcharged on my last invoice.",61            "This charge should not appear on my account.",62            "I dispute this billing error.",63            "The pricing on this invoice is incorrect.",64        ],65        # Refund requests66        "Billing-Refund": [67            "I would like to request a refund.",68            "Please process a refund for my order.",69            "When will my refund be processed?",70            "I need to return and get refunded.",71        ],72        # Basic support73        "Support": [74            "I need technical assistance with your product.",75            "How can I contact customer support?",76            "I have a question about your services.",77            "Can you help me with this issue?",78        ],79        # EVERYTHING'S ON FIRE type support80        "Support-Urgent": [81            "URGENT: My system is down and I need immediate help!",82            "CRITICAL: This is affecting production!",83            "Emergency support needed immediately!",84            "This is urgent and needs immediate attention!",85        ],86        # Normal support that can wait87        "Support-Normal": [88            "I am having trouble with a basic feature.",89            "Can you walk me through the process?",90            "I need help understanding this feature.",91            "How do I use this functionality?",92        ],93        # Technical stuff94        "Technical": [95            "I am experiencing a technical issue with your software.",96            "There seems to be a bug in the system.",97            "Technical support needed for integration.",98            "I need help with API documentation.",99        ],100        # Code broken101        "Technical-Bug": [102            "I found a bug in the system that crashes the app.",103            "The system throws an error when I try to login.",104            "Error: Function X is not working correctly.",105            "Bug report: System crashes on this operation.",106        ],107        # New features people want108        "Technical-Feature": [109            "I would like to request a new feature.",110            "Can you add functionality for X?",111            "Feature request: Please add support for this.",112            "Would it be possible to implement this capability?",113        ],114        # HR stuff115        "HR": [116            "I have an HR-related question.",117            "I need assistance with HR matters.",118            "Can you help with HR policies?",119            "I have a question about HR.",120        ],121        # Money questions122        "HR-Payroll": [123            "I have a question about my paycheck.",124            "My salary appears to be incorrect.",125            "When is the next payroll cycle?",126            "Can you explain my payroll deductions?",127        ],128        # Health insurance etc129        "HR-Benefits": [130            "I need to enroll in health benefits.",131            "Can I update my benefits information?",132            "What benefits am I eligible for?",133            "I have questions about the benefits plan.",134        ],135        # Someone's upset136        "HR-Complaint": [137            "I would like to file a formal complaint.",138            "I need to report a workplace issue.",139            "This is a serious HR matter that needs attention.",140            "I am filing an official complaint.",141        ],142        # Legal stuff143        "Legal": [144            "I need legal advice regarding a contract.",145            "I have legal questions about terms.",146            "Can you help with legal documentation?",147            "I need assistance with legal matters.",148        ],149        # Contracts150        "Legal-Contract": [151            "I need review of this contract before signing.",152            "The contract terms need clarification.",153            "I have questions about contract obligations.",154            "Can you explain these contract provisions?",155        ],156        # Compliance - boring but important157        "Legal-Compliance": [158            "We need to ensure compliance with regulations.",159            "I have questions about legal compliance.",160            "Are we meeting compliance requirements?",161            "I need guidance on regulatory compliance.",162        ],163        # Executive stuff164        "Executive": [165            "I would like to schedule a meeting with management.",166            "This is an executive-level inquiry.",167            "Can I speak with a senior manager?",168            "I need executive attention for this matter.",169        ],170        # Big picture strategy171        "Executive-Strategic": [172            "Strategic partnership opportunity for discussion.",173            "I would like to discuss corporate strategy.",174            "This requires strategic-level decision making.",175            "Executive strategic initiative to discuss.",176        ],177        # Finance178        "Finance": [179            "I have questions about financial matters.",180            "Can you provide financial statements?",181            "I need to review financial records.",182            "Financial assistance or information needed.",183        ],184        # Marketing185        "Marketing": [186            "I am interested in marketing partnerships.",187            "Can we discuss marketing opportunities?",188            "I have a marketing proposal for you.",189            "Marketing collaboration inquiry.",190        ],191        # Operations192        "Operations": [193            "I have operational questions.",194            "Operations scheduling needs adjustment.",195            "Operational efficiency improvement proposal.",196            "I need to discuss operational procedures.",197        ],198    }199    200    def __init__(self, difficulty: str, seed: int = None):201        self.difficulty = difficulty202        self.config = get_task_config(difficulty)203        if seed is not None:204            np.random.seed(seed)205            random.seed(seed)206        207        # Generate category list based on difficulty208        if difficulty == "easy":209            self.categories = ["General", "Billing", "Support", "Technical", "HR"]210        elif difficulty == "medium":211            self.categories = [212                "General", "Billing", "Billing-Dispute", "Support",213                "Technical", "Technical-Bug", "HR-Payroll", "HR-Benefits",214                "Legal", "Executive"215            ]216        else:  # hard217            self.categories = list(self.DOCUMENT_TEMPLATES.keys())218        219        # Initialize TF-IDF vectorizer220        all_docs = []221        for templates in self.DOCUMENT_TEMPLATES.values():222            all_docs.extend(templates)223        224        self.vectorizer = TfidfVectorizer(225            max_features=self.config["feature_dim"],226            lowercase=True,227            stop_words='english'228        )229        self.vectorizer.fit(all_docs)230    231    def generate_task_data(self) -> Tuple[List[Dict], np.ndarray]:232        """233        Create a batch of documents.234        Pick random categories, grab templates, add some noise, extract features.235        """236        documents = []237        labels = []238        239        num_docs = self.config["num_documents"]240        241        for i in range(num_docs):242            # Pick a random category for this doc243            category = random.choice(self.categories)244            category_idx = self.categories.index(category)245            246            # Get a template from this category247            template = random.choice(self.DOCUMENT_TEMPLATES[category])248            249            # Add some variation so it's not too repetitive250            variations = self._add_variation(template)251            252            # Extract features (TF-IDF)253            features = self.vectorizer.transform([variations]).toarray()[0]254            255            # Normalize to [-1, 1] range256            features = 2 * features / (np.max(features) + 1e-8) - 1257            258            # Create document dict259            doc = {260                "id": f"doc_{i:06d}",261                "content": variations,262                "word_count": len(variations.split()),263                "has_urgency_markers": self._has_urgency(variations),264                "features": features.tolist(),265                "true_category": category,266            }267            268            documents.append(doc)269            labels.append(category_idx)270        271        return documents, np.array(labels)272    273    def _add_variation(self, template: str) -> str:274        """Add some noise/variation to make it feel less templated"""275        words = template.split()276        277        # Sometimes add filler words278        if random.random() < 0.3:279            filler = random.choice([280                "I think", "Regarding", "In addition", "Furthermore",281                "Also", "Additionally", "Moreover", "Please note"282            ])283            words.insert(random.randint(0, len(words)), filler)284        285        # Maybe add some detail286        if random.random() < 0.2:287            details = random.choice([288                " regarding order #12345.",289                " It's urgent.",290                " This needs immediate attention.",291                " Thank you for your help.",292                " Please let me know ASAP.",293            ])294            words.append(details)295        296        return " ".join(words)297    298    def _has_urgency(self, text: str) -> bool:299        """Check if text sounds urgent/important"""300        urgency_words = ["urgent", "emergency", "critical", "immediate", "asap", "urgent!"]301        text_lower = text.lower()302        return any(word in text_lower for word in urgency_words)303 304 305# For testing306if __name__ == "__main__":307    # Test data generation308    for difficulty in ["easy", "medium", "hard"]:309        gen = TaskDataGenerator(difficulty, seed=42)310        docs, labels = gen.generate_task_data()311        print(f"\n{difficulty.upper()} Task:")312        print(f"  Documents: {len(docs)}")313        print(f"  Categories: {len(set(labels))}")314        print(f"  Sample doc: {docs[0]['id']}")315        print(f"  Sample content: {docs[0]['content'][:50]}...")316        print(f"  Features shape: {len(docs[0]['features'])}")317