xfmr-zem 0.2.0__py3-none-any.whl

xfmr_zem/server.py

@@ -0,0 +1,188 @@
+
+from typing import Any, Callable, Dict, List, Optional, Union
+import yaml
+from pathlib import Path
+from fastmcp import FastMCP
+import inspect
+
+class ZemServer(FastMCP):
+    """
+    Base class for Zem MCP Servers.
+    Extends FastMCP to support parameter loading and standardized tool registration.
+    """
+    
+    def __init__(
+        self,
+        name: str,
+        Dependencies: Optional[List[str]] = None,
+        parameter_file: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(name=name, **kwargs)
+        self.parameter_file = parameter_file
+        self.parameters = {}
+        
+        # 1. Load from file
+        if parameter_file:
+            self.load_parameters(parameter_file)
+            
+        # 2. Override with env params (from PipelineClient)
+        import os
+        env_params_str = os.environ.get("ZEM_PARAMETERS")
+        if env_params_str:
+            try:
+                env_params = yaml.safe_load(env_params_str)
+                if isinstance(env_params, dict):
+                    self._merge_parameters(env_params)
+            except Exception as e:
+                print(f"Error loading ZEM_PARAMETERS: {e}")
+
+    def load_parameters(self, file_path: str) -> Dict[str, Any]:
+        """Load parameters from YAML file and merge them."""
+        path = Path(file_path)
+        if path.exists():
+            with open(path, "r") as f:
+                file_params = yaml.safe_load(f) or {}
+                self._merge_parameters(file_params)
+            return self.parameters
+        return {}
+
+    def _merge_parameters(self, new_params: Dict[str, Any]):
+        """Deep merge and dot-notation expansion for parameters."""
+        for key, value in new_params.items():
+            if "." in key:
+                # Expand "tool.param" to {"tool": {"param": value}}
+                parts = key.split(".")
+                d = self.parameters
+                for part in parts[:-1]:
+                    if part not in d or not isinstance(d[part], dict):
+                        d[part] = {}
+                    d = d[part]
+                
+                last_part = parts[-1]
+                if isinstance(value, dict) and last_part in d and isinstance(d[last_part], dict):
+                    self._deep_update(d[last_part], value)
+                else:
+                    d[last_part] = value
+            else:
+                # Top level merge
+                if isinstance(value, dict) and key in self.parameters and isinstance(self.parameters[key], dict):
+                    self._deep_update(self.parameters[key], value)
+                else:
+                    self.parameters[key] = value
+
+    def _deep_update(self, target: Dict[str, Any], source: Dict[str, Any]):
+        """Helper for deep dictionary update."""
+        for k, v in source.items():
+            if isinstance(v, dict) and k in target and isinstance(target[k], dict):
+                self._deep_update(target[k], v)
+            else:
+                target[k] = v
+
+    # Removed custom tool decorator to fix multiple values for argument 'name' error
+    # Inherit directly from FastMCP.tool
+
+    def get_data(self, data: Any) -> List[Dict[str, Any]]:
+        """
+        Standardized way to get data, supporting both direct lists and file references.
+        """
+        import os
+        from loguru import logger
+        
+        logger.debug(f"Server {self.name}.get_data input type: {type(data)}")
+
+        if isinstance(data, list):
+            return data
+        
+        if isinstance(data, dict):
+            if "path" in data:
+                path = data["path"]
+                ext = os.path.splitext(path)[1].lower()
+                
+                logger.debug(f"Server {self.name}: Loading reference {path}")
+                if ext == ".jsonl":
+                    import json
+                    with open(path, "r", encoding="utf-8") as f:
+                        return [json.loads(line) for line in f if line.strip()]
+                elif ext == ".csv":
+                    import pandas as pd
+                    return pd.read_csv(path).to_dict(orient="records")
+                elif ext == ".parquet":
+                    import pandas as pd
+                    return pd.read_parquet(path).to_dict(orient="records")
+            else:
+                # Single record dictionary
+                logger.debug(f"Server {self.name}: Wrapping single dict in list")
+                return [data]
+        
+        if isinstance(data, str):
+            # 1. URL/Cloud URI or Local Path
+            is_uri = data.startswith(("http", "s3://", "gs://"))
+            if is_uri or os.path.exists(data):
+                ext = os.path.splitext(data)[1].lower()
+                logger.debug(f"Server {self.name}: Loading data from {data}")
+                try:
+                    import pandas as pd
+                    if ext == ".parquet":
+                        return pd.read_parquet(data).to_dict(orient="records")
+                    elif ext == ".csv":
+                        return pd.read_csv(data).to_dict(orient="records")
+                    elif ext == ".jsonl":
+                        if is_uri:
+                            return pd.read_json(data, lines=True).to_dict(orient="records")
+                        import json
+                        with open(data, "r", encoding="utf-8") as f:
+                            return [json.loads(line) for line in f if line.strip()]
+                except Exception as e:
+                    logger.error(f"Error loading data from {data}: {e}")
+            
+            # 2. Treat as raw text or JSON string
+            try:
+                import json
+                parsed = json.loads(data)
+                if isinstance(parsed, list): return parsed
+                if isinstance(parsed, dict): return [parsed]
+            except:
+                pass
+            return [{"text": data}]
+            
+        logger.warning(f"Server {self.name}: Unrecognized data type {type(data)}")
+        return [{"raw": str(data)}]
+
+    def save_output(self, data: Any, format: str = "parquet") -> Dict[str, Any]:
+        """
+        Saves output to a temporary file and returns a reference.
+        Prevents large data from being sent over JSON-RPC.
+        """
+        import uuid
+        import os
+        from loguru import logger
+        
+        base_dir = "/tmp/zem_artifacts"
+        os.makedirs(base_dir, exist_ok=True)
+        
+        file_id = str(uuid.uuid4())[:8]
+        path = os.path.join(base_dir, f"{self.name}_output_{file_id}.{format}")
+        
+        logger.info(f"Server {self.name}: Saving result to reference {path}")
+        
+        if format == "parquet":
+            import pandas as pd
+            logger.debug(f"Server {self.name}: Converting to DataFrame, data type: {type(data)}")
+            try:
+                df = pd.DataFrame(data)
+                df.to_parquet(path, index=False)
+            except Exception as e:
+                logger.error(f"Server {self.name}: Failed to create DataFrame from {type(data)}: {e}")
+                raise
+        elif format == "jsonl":
+            import json
+            with open(path, "w", encoding="utf-8") as f:
+                for item in data:
+                    f.write(json.dumps(item, ensure_ascii=False) + "\n")
+        
+        return {"path": path, "type": format, "size": os.path.getsize(path)}
+
+    def run(self, transport: str = "stdio"):
+        """Run the server."""
+        super().run(transport=transport, show_banner=False)

xfmr_zem/servers/data_juicer/parameter.yaml

@@ -0,0 +1,17 @@
+# Technical Parameters for DataJuicer Server
+
+clean_content:
+  remove_html: true
+  remove_emojis: false
+  text_column: "text"
+
+refining_filter:
+  min_len: 10
+  max_len: 100000
+  alphanumeric_ratio: 0.1
+  text_column: "text"
+
+language_id:
+  expected_lang: "vi"
+  min_score: 0.8
+  text_column: "text"

xfmr_zem/servers/data_juicer/server.py

@@ -0,0 +1,113 @@
+import os
+import sys
+import re
+from typing import Any, Dict, List, Optional
+from xfmr_zem.server import ZemServer
+from loguru import logger
+
+# Setup logging
+logger.remove()
+logger.add(sys.stderr, level="INFO")
+
+server = ZemServer("data_juicer", parameter_file=os.path.join(os.path.dirname(__file__), "parameter.yaml"))
+
+@server.tool()
+def clean_content(
+    data: Any,
+    remove_html: bool = True,
+    remove_emojis: bool = False,
+    text_column: str = "text"
+) -> Any:
+    """
+    Flexible content cleaning tool using DataJuicer logic.
+    """
+    items = server.get_data(data)
+    if not items: return []
+    
+    logger.info(f"DataJuicer: Cleaning content (remove_html={remove_html}, remove_emojis={remove_emojis})")
+    
+    for item in items:
+        if text_column not in item: continue
+        text = str(item[text_column])
+        
+        if remove_html:
+            text = re.sub(r'<[^>]+>', '', text)
+            
+        if remove_emojis:
+            # Targeted emoji removal: remove high-plane Unicode characters (likely emojis) 
+            # while preserving Vietnamese diacritics and other standard Unicode text.
+            text = re.sub(r'[\U00010000-\U0010ffff]', '', text)
+            
+        item[text_column] = text.strip()
+        
+    return server.save_output(items)
+
+@server.tool()
+def refining_filter(
+    data: Any,
+    min_len: int = 10,
+    max_len: int = 100000,
+    alphanumeric_ratio: float = 0.1,
+    text_column: str = "text"
+) -> Any:
+    """
+    Filter items based on technical refining metrics.
+    """
+    items = server.get_data(data)
+    if not items: return []
+    
+    logger.info(f"DataJuicer: Refining filter (min_len={min_len}, alpha_ratio={alphanumeric_ratio})")
+    
+    filtered = []
+    for item in items:
+        text = str(item.get(text_column, ""))
+        text_len = len(text)
+        
+        if not (min_len <= text_len <= max_len):
+            continue
+            
+        # Alphanumeric ratio check
+        alnum_count = len([c for c in text if c.isalnum()])
+        if text_len > 0 and (alnum_count / text_len) < alphanumeric_ratio:
+            continue
+            
+        filtered.append(item)
+        
+    return server.save_output(filtered)
+
+@server.tool()
+def language_id(
+    data: Any,
+    expected_lang: str = "vi",
+    min_score: float = 0.8,
+    text_column: str = "text"
+) -> Any:
+    """
+    Heuristic language identification tool.
+    """
+    items = server.get_data(data)
+    if not items: return []
+    
+    logger.info(f"DataJuicer: Filtering for language '{expected_lang}'")
+    
+    # Heuristic for Vietnamese
+    vi_keywords = {'và', 'của', 'là', 'có', 'được', 'trong', 'cho', 'này', 'với', 'các'}
+    
+    filtered = []
+    for item in items:
+        text = str(item.get(text_column, "")).lower()
+        words = set(text.split())
+        
+        if expected_lang == "vi":
+            matches = len(words & vi_keywords)
+            # Heuristic score: ratio of keywords found (simplified)
+            if matches >= 2:
+                filtered.append(item)
+        else:
+            # Fallback for other languages
+            filtered.append(item)
+            
+    return server.save_output(filtered)
+
+if __name__ == "__main__":
+    server.run()

xfmr_zem/servers/instruction_gen/parameter.yaml

@@ -0,0 +1,12 @@
+# Technical Parameters for Instruction Generation Server
+# These are default technical values. Override them in your pipeline YAML.
+
+generate_qa_pairs:
+  base_url: "http://localhost:8000/v1"
+  model: "default"
+  num_pairs: 3
+  text_column: "text"
+  prompt_template: "Dựa trên văn bản pháp luật sau, hãy tạo {num_pairs} cặp Câu hỏi và Trả lời chi tiết. Trả về định dạng JSON list: [{{'q': '...', 'a': '...'}}]. Văn bản: {text}"
+
+complexity_scorer:
+  text_column: "text"

xfmr_zem/servers/instruction_gen/server.py

@@ -0,0 +1,90 @@
+import os
+import sys
+import json
+import httpx
+from typing import Any, Dict, List, Optional
+from xfmr_zem.server import ZemServer
+from loguru import logger
+
+# Setup logging
+logger.remove()
+logger.add(sys.stderr, level="INFO")
+
+server = ZemServer("instruction", parameter_file=os.path.join(os.path.dirname(__file__), "parameter.yaml"))
+
+@server.tool()
+def generate_qa_pairs(
+    data: Any = None,
+    base_url: str = "http://localhost:8000/v1",
+    model: str = "default",
+    num_pairs: int = 3,
+    text_column: str = "text",
+    prompt_template: str = "Dựa trên văn bản pháp luật sau, hãy tạo {num_pairs} cặp Câu hỏi và Trả lời chi tiết. Trả về định dạng JSON list: [{{'q': '...', 'a': '...'}}]. Văn bản: {text}"
+) -> Any:
+    raw_data = server.get_data(data)
+    if not raw_data: return []
+    
+    # Handle wrapped data from previous steps
+    if isinstance(raw_data, dict) and 'data' in raw_data:
+        actual_items = raw_data['data']
+    else:
+        actual_items = raw_data
+    
+    logger.info(f"InstructionGen: Connecting to vLLM at {base_url} for {len(actual_items)} items")
+    
+    processed_items = []
+    with httpx.Client(timeout=60.0) as client:
+        for item in actual_items:
+            text = str(item.get(text_column, ""))
+            if not text: continue
+            
+            prompt = prompt_template.format(num_pairs=num_pairs, text=text)
+            try:
+                response = client.post(
+                    f"{base_url}/chat/completions",
+                    json={
+                        "model": model,
+                        "messages": [{"role": "user", "content": prompt}],
+                        "temperature": 0.3
+                    }
+                )
+                if response.status_code == 200:
+                    raw_content = response.json()["choices"][0]["message"]["content"]
+                    start = raw_content.find('[')
+                    end = raw_content.rfind(']') + 1
+                    if start != -1 and end != -1:
+                        item["instructions"] = json.loads(raw_content[start:end])
+                else:
+                    item["instructions"] = []
+            except Exception:
+                item["instructions"] = []
+            processed_items.append(item)
+            
+    return server.save_output(processed_items)
+
+@server.tool()
+def complexity_scorer(
+    data: Any = None, 
+    text_column: str = "text"
+) -> Any:
+    raw_data = server.get_data(data)
+    if not raw_data: return []
+    
+    if isinstance(raw_data, dict) and 'data' in raw_data:
+        actual_items = raw_data['data']
+    else:
+        actual_items = raw_data
+    
+    logger.info(f"InstructionGen: Scoring complexity for {len(actual_items)} items")
+    
+    for item in actual_items:
+        text = str(item.get(text_column, ""))
+        words = text.split()
+        unique_words = set(words)
+        score = min(1.0, (len(unique_words) / 100) * (len(words) / 500))
+        item["complexity_score"] = round(float(score), 2)
+        
+    return server.save_output(actual_items)
+
+if __name__ == "__main__":
+    server.run()

xfmr_zem/servers/io/parameter.yaml

@@ -0,0 +1,10 @@
+# IO Server Configuration
+# Default settings for loading and saving data
+
+load_settings:
+  chunk_size: 1000
+  encoding: "utf-8"
+
+save_settings:
+  overwrite: true
+  encoding: "utf-8"

xfmr_zem/servers/io/server.py

@@ -0,0 +1,95 @@
+import os
+import json
+import pandas as pd
+from typing import Any, Dict, List, Optional
+from xfmr_zem.server import ZemServer
+from loguru import logger
+import sys
+
+# Initialize server
+server = ZemServer("io", parameter_file=os.path.join(os.path.dirname(__file__), "parameter.yaml"))
+
+@server.tool()
+def load_jsonl(path: str, return_reference: bool = False) -> Any:
+    """
+    Load data from a JSONL file. If return_reference is True, returns metadata instead of data.
+    """
+    logger.info(f"IO: Loading JSONL from {path}")
+    if not os.path.exists(path):
+        raise FileNotFoundError(f"File not found: {path}")
+    
+    if return_reference:
+        return {"path": path, "type": "jsonl", "size": os.path.getsize(path)}
+
+    data = []
+    with open(path, "r", encoding="utf-8") as f:
+        for line in f:
+            if line.strip():
+                data.append(json.loads(line))
+    return data
+
+@server.tool()
+def save_jsonl(data: Any, path: str) -> Dict[str, Any]:
+    """Save data (List or Path reference) to JSONL."""
+    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+    
+    # If data is a reference (dict with path), we might want to move/copy it
+    # but for now, assume data is a List if we are in save_jsonl
+    if isinstance(data, dict) and "path" in data:
+        import shutil
+        shutil.copy(data["path"], path)
+        return {"status": "copied", "path": path}
+
+    with open(path, "w", encoding="utf-8") as f:
+        for item in data:
+            f.write(json.dumps(item, ensure_ascii=False) + "\n")
+    return {"status": "success", "path": path, "count": len(data)}
+
+@server.tool()
+def load_parquet(path: str, return_reference: bool = True) -> Any:
+    """
+    Load data from a Parquet file. Defaults to return_reference=True for big data.
+    """
+    logger.info(f"IO: Handling Parquet from {path}")
+    if not os.path.exists(path):
+         raise FileNotFoundError(f"File not found: {path}")
+
+    if return_reference:
+        return {"path": path, "type": "parquet", "size": os.path.getsize(path)}
+    
+    df = pd.read_parquet(path)
+    return df.to_dict(orient="records")
+
+@server.tool()
+def save_parquet(data: Any, path: str) -> Dict[str, Any]:
+    """
+    Save data to a Parquet file.
+    """
+    logger.info(f"IO: Saving Parquet to {path}")
+    os.makedirs(os.path.dirname(os.path.abspath(path)), exist_ok=True)
+    
+    if isinstance(data, dict) and "path" in data:
+        src = data["path"]
+        logger.info(f"IO: Copying reference from {src} to {path}")
+        import shutil
+        shutil.copy(src, path)
+        return {"status": "copied", "path": path}
+    
+    logger.info("IO: Saving raw data to Parquet")
+    df = pd.DataFrame(data)
+    df.to_parquet(path, index=False)
+    return {"status": "success", "path": path, "count": len(data)}
+
+@server.tool()
+def scan_directory(path: str, pattern: str = "*") -> List[Dict[str, Any]]:
+    """
+    Scan a directory and return a list of file references.
+    Useful for parallelizing big data processing.
+    """
+    import glob
+    logger.info(f"IO: Scanning directory {path} with pattern {pattern}")
+    files = glob.glob(os.path.join(path, pattern))
+    return [{"path": f, "type": "file_ref"} for f in files]
+
+if __name__ == "__main__":
+    server.run()

xfmr_zem/servers/llm/server.py

@@ -0,0 +1,47 @@
+from xfmr_zem.server import ZemServer
+from typing import Any, List, Optional
+import os
+import json
+import requests
+
+mcp = ZemServer("LLM-Curation")
+
+def _call_llm(prompt: str, provider: str = "ollama", model: Optional[str] = None) -> str:
+    """Helper to route LLM calls."""
+    if provider == "ollama":
+        model = model or "llama3"
+        try:
+            url = os.environ.get("OLLAMA_URL", "http://localhost:11434/api/generate")
+            response = requests.post(url, json={"model": model, "prompt": prompt, "stream": False})
+            return response.json().get("response", "")
+        except Exception:
+            return f"[Ollama Error] Could not connect to {model}. Fallback: processed {prompt[:20]}..."
+    elif provider == "openai":
+        api_key = os.environ.get("OPENAI_API_KEY")
+        if not api_key: return "[OpenAI Error] Missing API Key."
+        # Placeholder for real openai client call
+        return f"[OpenAI Mock] Classification/Response for: {prompt[:20]}..."
+    return f"[Mock] {prompt}"
+
+@mcp.tool()
+def mask_pii(data: Any, provider: str = "ollama") -> List[Any]:
+    """Smart PII masking using LLM."""
+    dataset = mcp.get_data(data)
+    for item in dataset:
+        if "text" in item:
+            prompt = f"Remove all PII from this text and return ONLY the cleaned text: {item['text']}"
+            item["text"] = _call_llm(prompt, provider=provider)
+    return dataset
+
+@mcp.tool()
+def classify_domain(data: Any, categories: List[str] = ["Tech", "Finance", "Legal"], provider: str = "ollama") -> List[Any]:
+    """Classify data domain using LLM."""
+    dataset = mcp.get_data(data)
+    for item in dataset:
+        if "text" in item:
+            prompt = f"Classify this text into one of {categories}. Return ONLY the category name: {item['text']}"
+            item["domain"] = _call_llm(prompt, provider=provider).strip()
+    return dataset
+
+if __name__ == "__main__":
+    mcp.run()

xfmr_zem/servers/nemo_curator/parameter.yaml

@@ -0,0 +1,17 @@
+# Technical Parameters for Nemo Curator Server
+# These are default technical values. Override them in your pipeline YAML.
+
+normalize:
+  normalization: "NFC"
+  text_column: "text"
+  cleanup_patterns: []
+
+quality_filter:
+  min_words: 50
+  max_non_alpha_ratio: 0.25
+  text_column: "text"
+
+deduplicate:
+  text_column: "text"
+  method: "exact"
+  threshold: 0.85