Create processor.py

processor.py

@@ -0,0 +1,168 @@
+import json
+import logging
+from datasets import load_dataset, Dataset, Features, Value
+from huggingface_hub import HfApi
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class DatasetCommandCenter:
+    def __init__(self, token=None):
+        self.token = token
+
+    def inspect_dataset(self, dataset_id, split="train", config=None):
+        """
+        Loads the first N rows and detects JSON structures.
+        """
+        try:
+            # Load in streaming mode to avoid RAM issues
+            ds_stream = load_dataset(dataset_id, config, split=split, streaming=True, token=self.token)
+            
+            # Peek at the first 5 rows
+            sample_rows = []
+            for i, row in enumerate(ds_stream):
+                if i >= 5: break
+                sample_rows.append(row)
+
+            # Analyze Columns
+            analysis = {}
+            if sample_rows:
+                keys = sample_rows[0].keys()
+                for k in keys:
+                    # Check if it looks like JSON
+                    is_json_candidate = False
+                    sample_val = sample_rows[0][k]
+                    
+                    if isinstance(sample_val, str):
+                        sample_val = sample_val.strip()
+                        if (sample_val.startswith('{') and sample_val.endswith('}')) or \
+                           (sample_val.startswith('[') and sample_val.endswith(']')):
+                            try:
+                                json.loads(sample_val)
+                                is_json_candidate = True
+                            except:
+                                pass
+                    
+                    analysis[k] = {
+                        "type": str(type(sample_rows[0][k]).__name__),
+                        "is_json": is_json_candidate
+                    }
+
+            return {
+                "status": "success",
+                "samples": sample_rows,
+                "analysis": analysis,
+                "dataset_id": dataset_id
+            }
+        except Exception as e:
+            return {"status": "error", "message": str(e)}
+
+    def _apply_transformations(self, row, recipe):
+        """
+        Applies a single row transformation based on the recipe.
+        Recipe format:
+        {
+            "json_expansions": [{"col": "meta", "keys": ["url", "id"]}],
+            "renames": {"old_col": "new_col"},
+            "drops": ["unwanted_col"],
+            "filters": ["len(text) > 50"]  # List of python eval strings
+        }
+        """
+        new_row = row.copy()
+        
+        # 1. JSON Expansions
+        if "json_expansions" in recipe:
+            for item in recipe["json_expansions"]:
+                col_name = item["col"]
+                target_keys = item["keys"] # List of keys to extract
+                
+                if col_name in new_row and isinstance(new_row[col_name], str):
+                    try:
+                        data = json.loads(new_row[col_name])
+                        for key in target_keys:
+                            # Support simple dot notation if needed, currently 1 level
+                            clean_key = key.replace('.', '_')
+                            new_col_name = f"{col_name}_{clean_key}"
+                            new_row[new_col_name] = data.get(key, None)
+                    except:
+                        # If parsing fails, fill None
+                        for key in target_keys:
+                            new_row[f"{col_name}_{key}"] = None
+
+        # 2. Renames
+        if "renames" in recipe:
+            for old, new in recipe["renames"].items():
+                if old in new_row:
+                    new_row[new] = new_row.pop(old)
+
+        # 3. Drops
+        if "drops" in recipe:
+            for drop_col in recipe["drops"]:
+                if drop_col in new_row:
+                    del new_row[drop_col]
+
+        return new_row
+
+    def _passes_filter(self, row, filters):
+        """
+        Safe-ish eval for filtering.
+        """
+        if not filters:
+            return True
+        
+        # We create a local context with the row data accessible as variables
+        # e.g., if row has 'text', user can write "len(text) > 5"
+        context = row.copy()
+        
+        for f_str in filters:
+            try:
+                if not eval(f_str, {}, context):
+                    return False
+            except Exception as e:
+                # If filter crashes (e.g. missing column), we skip the row or default to False
+                return False
+        return True
+
+    def preview_transform(self, dataset_id, split, recipe):
+        """Return a transformed sample of 5 rows"""
+        ds_stream = load_dataset(dataset_id, split=split, streaming=True, token=self.token)
+        processed = []
+        
+        for row in ds_stream:
+            if len(processed) >= 5: break
+            
+            # Apply Filter
+            if not self._passes_filter(row, recipe.get("filters", [])):
+                continue
+
+            # Apply Transform
+            trans_row = self._apply_transformations(row, recipe)
+            processed.append(trans_row)
+            
+        return processed
+
+    def process_and_push(self, source_id, split, target_id, recipe, max_rows=None):
+        """
+        The heavy lifter: Streams, Transforms, Filters, and Pushes to Hub.
+        """
+        logger.info(f"Starting job: {source_id} -> {target_id}")
+        
+        def gen():
+            ds_stream = load_dataset(source_id, split=split, streaming=True, token=self.token)
+            count = 0
+            for row in ds_stream:
+                if max_rows and count >= int(max_rows):
+                    break
+                
+                if self._passes_filter(row, recipe.get("filters", [])):
+                    yield self._apply_transformations(row, recipe)
+                    count += 1
+
+        # Create new dataset from generator
+        # We let HF infer the features (schema) automatically from the first batch
+        new_dataset = Dataset.from_generator(gen)
+        
+        # Push
+        new_dataset.push_to_hub(target_id, token=self.token)
+        return {"status": "success", "rows_processed": len(new_dataset)}