src/datasets/curator.py

import sqlite3
import json
import os
import uuid
from typing import List, Dict, Any
from src.ontology.models import OntologyRecord

class DatasetCurator:
    def __init__(self, db_path: str, output_dir: str):
        self.db_path = db_path
        self.output_dir = output_dir
        self.nsfw_terms = set()
        os.makedirs(output_dir, exist_ok=True)

    def load_nsfw_terms(self):
        """Load NSFW terms from the database to use for classification."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # Load from special_all and special_s_e
        for table in ["special_all", "special_s_e"]:
            cursor.execute(f"SELECT value FROM {table}")
            for row in cursor.fetchall():
                term = row[0].strip().lower()
                if term:
                    self.nsfw_terms.add(term)
        
        conn.close()
        print(f"Loaded {len(self.nsfw_terms)} NSFW terms.")

    def is_nsfw(self, text: str) -> bool:
        """Check if a string contains NSFW terms using word boundaries."""
        if not text:
            return False
        text_lower = text.lower()
        
        # For performance, we can do a quick word-set check first for single-word terms
        import re
        words = set(re.findall(r'\w+', text_lower))
        if not words.isdisjoint(self.nsfw_terms):
            return True
            
        # Then check for multi-word terms (if any)
        # For now, most seem to be single words or joined words.
        # If we have multi-word terms, we'd need re.search with \b
        for term in self.nsfw_terms:
            if " " in term:
                if re.search(rf"\b{re.escape(term)}\b", text_lower):
                    return True
        return False

    def process_table(self, table_name: str, category: str, canonical_col: str = "value"):
        """Process a standard table with 'id' and 'value' columns."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        cursor.execute(f"SELECT * FROM {table_name}")
        rows = cursor.fetchall()
        
        # Get column names
        cursor.execute(f"PRAGMA table_info({table_name})")
        cols = [col[1] for col in cursor.fetchall()]
        
        records = []
        for row in rows:
            data = dict(zip(cols, row))
            canonical = data[canonical_col]
            
            # Basic normalization: strip and lowercase
            canonical = canonical.strip().lower()
            
            # Create OntologyRecord
            record = OntologyRecord(
                id=str(uuid.uuid4()),
                canonical=canonical,
                aliases=[], # We'll populate aliases later if needed
                category=category,
                source=f"db:{table_name}",
                nsfw=self.is_nsfw(canonical),
                synthetic=False,
                confidence=1.0
            )
            records.append(record.model_dump())
            
        conn.close()
        return records

    def process_characters(self):
        """Process the characters table which has a different schema."""
        conn = sqlite3.connect(self.db_path)
        cursor = conn.cursor()
        
        # We might want to join with franchises
        query = """
        SELECT c.id, c.name, c.core_tags, f.name as franchise
        FROM characters c
        LEFT JOIN franchises f ON c.franchise_id = f.id
        """
        cursor.execute(query)
        rows = cursor.fetchall()
        
        records = []
        for row in rows:
            cid, name, core_tags, franchise = row
            
            canonical = name
            aliases = []
            
            # Clean up canonical name: "Name from Franchise" -> "Name"
            # or "Name (Variant)" -> "Name"
            clean_name = canonical
            if " from " in clean_name:
                clean_name = clean_name.split(" from ")[0].strip()
            if " (" in clean_name:
                # Extract variant and add it to attributes or description later
                # For now, just get the base name
                clean_name = clean_name.split(" (")[0].strip()
            
            if clean_name != canonical:
                aliases.append(canonical) # Add the original as alias
                canonical = clean_name
            
            tags = []
            if core_tags:
                tags = [t.strip() for t in core_tags.split(",")]
            
            # Ensure unique aliases and exclude canonical
            aliases = list(set(a for a in aliases if a.lower() != canonical.lower()))
            
            description = f"Character from {franchise}" if franchise else ""
            
            metadata = {
                "franchise": franchise,
                "default_attributes": tags
            }
            
            record = OntologyRecord(
                id=str(uuid.uuid4()),
                canonical=canonical,
                aliases=aliases,
                category="character",
                description=description,
                tags=tags,
                source="db:characters",
                nsfw=self.is_nsfw(canonical) or self.is_nsfw(" ".join(aliases)),
                synthetic=False,
                confidence=1.0,
                metadata=metadata
            )
            records.append(record.model_dump())

            
        conn.close()
        return records

    def run(self):
        self.load_nsfw_terms()
        
        datasets = {
            "characters.json": self.process_characters(),
            "clothing.json": self.process_table("outfit", "clothing"),
            "hairstyles.json": self.process_table("hairstyle", "hairstyle"),
            "hair_colors.json": self.process_table("hair_color", "hair_color"),
            "eye_colors.json": self.process_table("eyes", "eye_color"),
            "scenes.json": self.process_table("scenario", "scene"),
            "emotions.json": self.process_table("emotion", "emotion"),
            "poses.json": self.process_table("pose", "pose"),
            "accessories.json": self.process_table("extras", "accessory"),
            "lighting.json": self.process_table("lighting", "lighting"),
            "styles.json": self.process_table("style", "style"),
            "effects.json": self.process_table("special_elements", "effect")
        }
        
        # Also process the special tables as entries themselves
        datasets["nsfw_all.json"] = self.process_table("special_all", "special", "value")
        datasets["nsfw_s_e.json"] = self.process_table("special_s_e", "special", "value")
        
        # Save all
        stats = {}
        for filename, data in datasets.items():
            file_path = os.path.join(self.output_dir, filename)
            with open(file_path, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
            stats[filename] = len(data)
            
        print("Dataset processing complete.")
        print(json.dumps(stats, indent=2))

if __name__ == "__main__":
    curator = DatasetCurator("fine_prompt_sdxl.db", "data/ontology")
    curator.run()