ritvik360/nl2sql-bench

1import json2import re3import sys4from pathlib import Path5 6sys.path.insert(0, str(Path(__file__).parent))7from data_factory.templates import ALL_TEMPLATES8 9# Define strict categorical swaps based on the exact schemas10SWAP_RULES = {11    "ecommerce": [12        (r"'gold'", r"gold", ["'silver'", "'bronze'"], ["silver", "bronze"]),13        (r"'delivered'", r"delivered", ["'pending'", "'processing'", "'shipped'", "'cancelled'"], ["pending", "processing", "shipped", "cancelled"]),14        (r"'India'", r"India", ["'USA'", "'Germany'", "'UK'", "'Canada'"], ["USA", "Germany", "UK", "Canada"])15    ],16    "healthcare": [17        (r"'severe'", r"severe", ["'mild'", "'moderate'"], ["mild", "moderate"]),18        (r"'completed'", r"completed", ["'scheduled'", "'cancelled'", "'no_show'"], ["scheduled", "cancelled", "no-show"])19    ],20    "finance": [21        (r"'active'", r"active", ["'dormant'", "'closed'"], ["dormant", "closed"]),22        (r"'credit'", r"credit", ["'debit'"], ["debit"]),23        (r"'verified'", r"verified", ["'pending'", "'rejected'"], ["pending", "rejected"])24    ],25    "hr": [26        (r"'active'", r"active", ["'resigned'", "'terminated'"], ["resigned", "terminated"])27    ]28}29 30def generate_swaps():31    expanded_templates = []32    33    for template in ALL_TEMPLATES:34        expanded_templates.append(template) # Keep the original35        domain = template["domain"]36        37        if domain not in SWAP_RULES:38            continue39            40        for sql_target, nl_target, sql_replacements, nl_replacements in SWAP_RULES[domain]:41            if re.search(sql_target, template["sql"], re.IGNORECASE):42                for sql_repl, nl_repl in zip(sql_replacements, nl_replacements):43                    new_template = template.copy()44                    45                    # Swap in SQL46                    new_template["sql"] = re.sub(sql_target, sql_repl, template["sql"], flags=re.IGNORECASE)47                    48                    # Swap in NL and Description49                    new_template["base_nl"] = re.sub(nl_target, nl_repl, template["base_nl"], flags=re.IGNORECASE)50                    new_template["description"] = re.sub(nl_target, nl_repl, template["description"], flags=re.IGNORECASE)51                    52                    # Create a unique ID53                    new_template["id"] = f"{template.get('id', 'temp')}_swap_{nl_repl.replace(' ', '_')}"54                    55                    expanded_templates.append(new_template)56                    57    return expanded_templates58 59if __name__ == "__main__":60    swapped = generate_swaps()61    print(f"Original Templates: {len(ALL_TEMPLATES)}")62    print(f"After Value Swapping: {len(swapped)}")63    64    with open("swapped_templates.json", "w") as f:65        json.dump(swapped, f, indent=2)66    print("Saved to swapped_templates.json")