sail_scripts/man/agent/data_compressor.py

"""
SAIL Data Compressor — Multi-Format Data Conversion Engine
============================================================
Converts raw data of ANY format into training-ready tensor representations.

Supported formats:
  • Text:        .txt, .md, .csv, .json → BPE tokenized tensors
  • Image:       .png, .jpg, .bmp, .webp → Resized + normalized feature tensors
  • Video:       .mp4, .avi, .mkv, .mov → Keyframe extraction → image pipeline
  • Document:    .pdf, .docx, .html → Text extraction → BPE tensors
  • Spreadsheet: .xlsx, .csv, .tsv → Structured column tensors

Each compressor returns a dict:
  {"type": "text|image|video|doc|sheet", "tensors": [...], "metadata": {...}}
"""

import os
import json
import hashlib
from pathlib import Path
from typing import Dict, List, Optional, Union
from datetime import datetime


# ─────────────────────────────────────────────────────────────────────────────
# Text Compressor
# ─────────────────────────────────────────────────────────────────────────────
class TextCompressor:
    """Compresses text files into BPE-tokenized packed tensors."""

    EXTENSIONS = {'.txt', '.md', '.csv', '.json', '.jsonl', '.log', '.xml'}

    def compress(self, file_path: str, tokenizer=None, max_length=2048) -> Dict:
        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
            text = f.read()

        # Basic cleaning
        text = text.strip()
        if not text:
            return {"type": "text", "tensors": [], "metadata": {"empty": True}}

        # Chunking for long texts
        chunks = self._chunk_text(text, max_length * 4)  # ~4 chars per token

        result = {
            "type": "text",
            "chunks": chunks,
            "metadata": {
                "source": file_path,
                "n_chunks": len(chunks),
                "total_chars": len(text),
                "hash": hashlib.md5(text[:1000].encode()).hexdigest()[:12],
            }
        }

        # Tokenize if tokenizer provided
        if tokenizer is not None:
            import torch
            token_tensors = []
            for chunk in chunks:
                ids = tokenizer.encode(chunk)
                if len(ids) > max_length:
                    ids = ids[:max_length]
                token_tensors.append(torch.tensor(ids, dtype=torch.long))
            result["tensors"] = token_tensors

        return result

    def _chunk_text(self, text: str, chunk_size: int) -> List[str]:
        if len(text) <= chunk_size:
            return [text]
        chunks = []
        for i in range(0, len(text), chunk_size - 200):  # 200 char overlap
            chunks.append(text[i:i + chunk_size])
        return chunks


# ─────────────────────────────────────────────────────────────────────────────
# Image Compressor
# ─────────────────────────────────────────────────────────────────────────────
class ImageCompressor:
    """Compresses images into normalized tensors, optionally with CLIP features."""

    EXTENSIONS = {'.png', '.jpg', '.jpeg', '.bmp', '.webp', '.tiff', '.gif'}
    DEFAULT_SIZE = (224, 224)

    def compress(self, file_path: str, size=None, use_clip=False) -> Dict:
        import torch
        import numpy as np

        size = size or self.DEFAULT_SIZE

        try:
            import cv2
            img = cv2.imread(file_path)
            if img is None:
                return {"type": "image", "tensors": [], "metadata": {"error": "unreadable"}}

            # Convert BGR → RGB
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            original_shape = img.shape

            # Resize
            img = cv2.resize(img, size, interpolation=cv2.INTER_LANCZOS4)

            # Normalize to [0, 1] then standard ImageNet normalization
            img = img.astype(np.float32) / 255.0
            mean = np.array([0.485, 0.456, 0.406])
            std = np.array([0.229, 0.224, 0.225])
            img = (img - mean) / std

            # Convert to tensor: (C, H, W)
            tensor = torch.from_numpy(img).permute(2, 0, 1).float()

            result = {
                "type": "image",
                "tensors": [tensor],
                "metadata": {
                    "source": file_path,
                    "original_shape": list(original_shape),
                    "compressed_shape": list(size),
                }
            }

            # Optional CLIP encoding for semantic features
            if use_clip:
                try:
                    from transformers import CLIPModel, CLIPProcessor
                    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
                    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
                    from PIL import Image
                    pil_img = Image.open(file_path).convert("RGB")
                    inputs = processor(images=pil_img, return_tensors="pt")
                    with torch.no_grad():
                        features = model.get_image_features(**inputs)
                    result["clip_features"] = features.squeeze(0)
                except Exception as e:
                    result["metadata"]["clip_error"] = str(e)

            return result

        except ImportError:
            return {"type": "image", "tensors": [], "metadata": {"error": "opencv not installed"}}


# ─────────────────────────────────────────────────────────────────────────────
# Video Compressor
# ─────────────────────────────────────────────────────────────────────────────
class VideoCompressor:
    """Extracts keyframes from video and compresses each through ImageCompressor."""

    EXTENSIONS = {'.mp4', '.avi', '.mkv', '.mov', '.webm', '.flv'}
    DEFAULT_FPS = 1  # Extract 1 frame per second

    def __init__(self):
        self.image_compressor = ImageCompressor()

    def compress(self, file_path: str, fps=None, max_frames=100) -> Dict:
        fps = fps or self.DEFAULT_FPS

        try:
            import cv2
            import torch
            import tempfile

            cap = cv2.VideoCapture(file_path)
            if not cap.isOpened():
                return {"type": "video", "tensors": [], "metadata": {"error": "cannot open"}}

            video_fps = cap.get(cv2.CAP_PROP_FPS)
            total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
            duration = total_frames / max(video_fps, 1)

            # Calculate frame interval
            frame_interval = max(1, int(video_fps / fps))

            frames = []
            frame_idx = 0
            while cap.isOpened() and len(frames) < max_frames:
                ret, frame = cap.read()
                if not ret:
                    break
                if frame_idx % frame_interval == 0:
                    # Save temp frame and compress
                    tmp_path = os.path.join(tempfile.gettempdir(), f"_sail_frame_{frame_idx}.jpg")
                    cv2.imwrite(tmp_path, frame)
                    result = self.image_compressor.compress(tmp_path)
                    if result["tensors"]:
                        frames.append(result["tensors"][0])
                    os.remove(tmp_path)
                frame_idx += 1

            cap.release()

            return {
                "type": "video",
                "tensors": frames,
                "metadata": {
                    "source": file_path,
                    "duration_sec": round(duration, 1),
                    "original_fps": video_fps,
                    "extracted_frames": len(frames),
                    "total_frames": total_frames,
                }
            }

        except ImportError:
            return {"type": "video", "tensors": [], "metadata": {"error": "opencv not installed"}}


# ─────────────────────────────────────────────────────────────────────────────
# Document Compressor
# ─────────────────────────────────────────────────────────────────────────────
class DocumentCompressor:
    """Extracts text from PDFs, DOCX, HTML and compresses via TextCompressor."""

    EXTENSIONS = {'.pdf', '.docx', '.doc', '.html', '.htm', '.rtf', '.epub'}

    def __init__(self):
        self.text_compressor = TextCompressor()

    def compress(self, file_path: str, tokenizer=None, max_length=2048) -> Dict:
        ext = Path(file_path).suffix.lower()
        text = ""

        if ext == '.pdf':
            text = self._extract_pdf(file_path)
        elif ext in ('.docx', '.doc'):
            text = self._extract_docx(file_path)
        elif ext in ('.html', '.htm'):
            text = self._extract_html(file_path)
        else:
            # Fallback: try reading as plain text
            try:
                with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                    text = f.read()
            except Exception:
                pass

        if not text.strip():
            return {"type": "document", "tensors": [], "metadata": {"error": "no text extracted"}}

        # Write extracted text to temp file and compress
        import tempfile
        tmp = os.path.join(tempfile.gettempdir(), "_sail_doc_extract.txt")
        with open(tmp, 'w', encoding='utf-8') as f:
            f.write(text)

        result = self.text_compressor.compress(tmp, tokenizer, max_length)
        result["type"] = "document"
        result["metadata"]["original_format"] = ext
        result["metadata"]["extracted_chars"] = len(text)
        os.remove(tmp)
        return result

    def _extract_pdf(self, path: str) -> str:
        try:
            from pypdf import PdfReader
            reader = PdfReader(path)
            return "\n".join(page.extract_text() or "" for page in reader.pages)
        except Exception as e:
            return f"[PDF extraction error: {e}]"

    def _extract_docx(self, path: str) -> str:
        try:
            import zipfile
            import xml.etree.ElementTree as ET
            with zipfile.ZipFile(path) as z:
                xml_content = z.read('word/document.xml')
            tree = ET.fromstring(xml_content)
            ns = {'w': 'http://schemas.openxmlformats.org/wordprocessingml/2006/main'}
            return "\n".join(
                node.text for node in tree.iter('{http://schemas.openxmlformats.org/wordprocessingml/2006/main}t')
                if node.text
            )
        except Exception as e:
            return f"[DOCX extraction error: {e}]"

    def _extract_html(self, path: str) -> str:
        try:
            import re
            with open(path, 'r', encoding='utf-8', errors='ignore') as f:
                html = f.read()
            # Strip tags
            text = re.sub(r'<[^>]+>', ' ', html)
            text = re.sub(r'&\w+;', ' ', text)
            return re.sub(r'\s+', ' ', text).strip()
        except Exception as e:
            return f"[HTML extraction error: {e}]"


# ─────────────────────────────────────────────────────────────────────────────
# Spreadsheet Compressor
# ─────────────────────────────────────────────────────────────────────────────
class SpreadsheetCompressor:
    """Converts spreadsheets into structured tensor representations."""

    EXTENSIONS = {'.xlsx', '.xls', '.csv', '.tsv'}

    def compress(self, file_path: str, max_rows=10000) -> Dict:
        import torch
        import numpy as np
        ext = Path(file_path).suffix.lower()

        try:
            import pandas as pd

            if ext in ('.csv', '.tsv'):
                sep = '\t' if ext == '.tsv' else ','
                df = pd.read_csv(file_path, sep=sep, nrows=max_rows)
            else:
                df = pd.read_excel(file_path, nrows=max_rows)

            # Separate numeric and text columns
            numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
            text_cols = df.select_dtypes(include=['object', 'string']).columns.tolist()

            tensors = []
            metadata = {
                "source": file_path,
                "shape": list(df.shape),
                "columns": list(df.columns),
                "numeric_cols": numeric_cols,
                "text_cols": text_cols,
            }

            # Numeric data → normalized tensor
            if numeric_cols:
                num_data = df[numeric_cols].fillna(0).values.astype(np.float32)
                # Z-score normalization per column
                mean = num_data.mean(axis=0, keepdims=True)
                std = num_data.std(axis=0, keepdims=True) + 1e-8
                num_data = (num_data - mean) / std
                tensors.append(torch.from_numpy(num_data))
                metadata["numeric_mean"] = mean.tolist()
                metadata["numeric_std"] = std.tolist()

            # Text data → concatenated string for tokenization
            if text_cols:
                text_concat = df[text_cols].fillna("").astype(str).apply(
                    lambda row: " | ".join(row), axis=1
                ).tolist()
                metadata["text_preview"] = text_concat[:3]
                metadata["text_rows"] = len(text_concat)
                # Store as raw text for later tokenization
                metadata["text_data"] = text_concat

            return {
                "type": "spreadsheet",
                "tensors": tensors,
                "metadata": metadata,
            }

        except ImportError:
            return {"type": "spreadsheet", "tensors": [], "metadata": {"error": "pandas not installed"}}


# ─────────────────────────────────────────────────────────────────────────────
# Main DataCompressor (Unified Interface)
# ─────────────────────────────────────────────────────────────────────────────
class DataCompressor:
    """
    Unified data compression engine.
    Auto-detects file type and routes to appropriate compressor.

    Usage:
        compressor = DataCompressor()
        result = compressor.compress("data/report.pdf")
        print(result["type"], len(result["tensors"]))

        # Batch compression
        results = compressor.compress_directory("data/raw/")
    """

    def __init__(self):
        self.text_compressor = TextCompressor()
        self.image_compressor = ImageCompressor()
        self.video_compressor = VideoCompressor()
        self.document_compressor = DocumentCompressor()
        self.spreadsheet_compressor = SpreadsheetCompressor()

        # Build extension → compressor map
        self._ext_map = {}
        for ext in TextCompressor.EXTENSIONS:
            self._ext_map[ext] = self.text_compressor
        for ext in ImageCompressor.EXTENSIONS:
            self._ext_map[ext] = self.image_compressor
        for ext in VideoCompressor.EXTENSIONS:
            self._ext_map[ext] = self.video_compressor
        for ext in DocumentCompressor.EXTENSIONS:
            self._ext_map[ext] = self.document_compressor
        for ext in SpreadsheetCompressor.EXTENSIONS:
            self._ext_map[ext] = self.spreadsheet_compressor

    def compress(self, file_path: str, **kwargs) -> Dict:
        """Auto-detect and compress a single file."""
        ext = Path(file_path).suffix.lower()
        compressor = self._ext_map.get(ext)

        if compressor is None:
            return {
                "type": "unknown",
                "tensors": [],
                "metadata": {"error": f"unsupported extension: {ext}"}
            }

        return compressor.compress(file_path, **kwargs)

    def compress_directory(self, dir_path: str, recursive=True, **kwargs) -> List[Dict]:
        """Compress all supported files in a directory."""
        results = []
        pattern = '**/*' if recursive else '*'

        for file_path in sorted(Path(dir_path).glob(pattern)):
            if not file_path.is_file():
                continue
            ext = file_path.suffix.lower()
            if ext not in self._ext_map:
                continue

            print(f"  Compressing: {file_path.name} ({ext})")
            try:
                result = self.compress(str(file_path), **kwargs)
                results.append(result)
            except Exception as e:
                print(f"  ✗ Error compressing {file_path.name}: {e}")
                results.append({
                    "type": "error",
                    "tensors": [],
                    "metadata": {"source": str(file_path), "error": str(e)}
                })

        print(f"\n  Compressed {len(results)} files from {dir_path}")
        return results

    def get_supported_extensions(self) -> Dict[str, List[str]]:
        """Returns all supported extensions grouped by type."""
        return {
            "text": sorted(TextCompressor.EXTENSIONS),
            "image": sorted(ImageCompressor.EXTENSIONS),
            "video": sorted(VideoCompressor.EXTENSIONS),
            "document": sorted(DocumentCompressor.EXTENSIONS),
            "spreadsheet": sorted(SpreadsheetCompressor.EXTENSIONS),
        }