feat: Add GraphNet Agent for automated sample extraction from HuggingFace models

Author: Dayuxiaoshui

graph_net/agent/README.md

@@ -0,0 +1,73 @@
+# GraphNet Agent
+
+自动样本抽取 Agent，实现从 HuggingFace ModelID 到 GraphNet Sample 的自动化转换。
+
+## 安装
+
+### 基础依赖
+```bash
+# 已包含在 GraphNet 主依赖中
+pip install torch torchvision
+```
+
+### Agent 可选依赖
+```bash
+# 安装 Agent 相关依赖（包括 huggingface_hub）
+pip install -e ".[agent]"
+
+# 或单独安装
+pip install huggingface_hub>=0.20.0
+```
+
+## 环境配置
+
+### 设置工作空间
+```bash
+export GRAPH_NET_EXTRACT_WORKSPACE=/path/to/your/workspace
+```
+
+或在代码中指定：
+```python
+from graph_net.agent import GraphNetAgent
+agent = GraphNetAgent(workspace="/path/to/workspace")
+```
+
+## 使用示例
+
+```python
+from graph_net.agent import GraphNetAgent
+
+# 初始化 Agent
+agent = GraphNetAgent(
+    workspace="./agent_workspace",
+    hf_token=None  # 可选，用于访问私有模型
+)
+
+# 运行提取
+success = agent.extract_sample("bert-base-uncased")
+
+if success:
+    print("✅ Sample extracted successfully")
+else:
+    print("❌ Extraction failed")
+```
+
+## 工作流程
+
+1. **Fetch**: 从 HuggingFace 下载模型
+2. **Analyze**: 解析 config.json 提取元数据
+3. **CodeGen**: 生成 run_model.py 脚本
+4. **Extract**: 执行脚本提取计算图
+5. **Deduplicate**: 检查是否与已有样本重复
+6. **Verify**: 验证样本完整性
+7. **Archive**: 保存 run_model.py 到样本目录
+
+## 测试
+
+```bash
+# 运行所有测试
+pytest graph_net/agent/tests/ -v
+
+# 运行实际模型测试（需要设置环境变量）
+TEST_REAL_RUN=1 pytest graph_net/agent/tests/test_real_run.py -v
+```

graph_net/agent/__init__.py

@@ -0,0 +1,10 @@
+"""
+GraphNet Auto Sample Extraction Agent
+
+This module provides an automated agent for extracting GraphNet samples
+from Hugging Face models.
+"""
+
+from graph_net.agent.graph_net_agent import GraphNetAgent
+
+__all__ = ["GraphNetAgent"]

graph_net/agent/code_generator/__init__.py

@@ -0,0 +1,6 @@
+"""Code generation modules"""
+
+from graph_net.agent.code_generator.base import BaseCodeGenerator
+from graph_net.agent.code_generator.template_generator import TemplateCodeGenerator
+
+__all__ = ["BaseCodeGenerator", "TemplateCodeGenerator"]

graph_net/agent/code_generator/base.py

@@ -0,0 +1,33 @@
+"""Base code generator interface"""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+from graph_net.agent.metadata_analyzer.model_metadata import ModelMetadata
+
+
+class BaseCodeGenerator(ABC):
+    """Base class for code generators"""
+    
+    @abstractmethod
+    def generate(
+        self,
+        model_dir: Path,
+        model_metadata: ModelMetadata,
+        output_dir: Path,
+    ) -> Path:
+        """
+        Generate run_model.py extraction script
+        
+        Args:
+            model_dir: Path to model directory
+            model_metadata: Model metadata extracted from configuration
+            output_dir: Output directory for generated script
+            
+        Returns:
+            Path to generated script file
+            
+        Raises:
+            CodeGenError: If code generation fails
+        """
+        pass

graph_net/agent/code_generator/template_generator.py

@@ -0,0 +1,173 @@
+"""Template-based code generator implementation"""
+
+from pathlib import Path
+from typing import Optional
+
+from graph_net.agent.metadata_analyzer.model_metadata import ModelMetadata
+from graph_net.agent.code_generator.base import BaseCodeGenerator
+from graph_net.agent.utils.exceptions import CodeGenError
+
+# Constants for safe vocab size calculation
+DEFAULT_VOCAB_SIZE = 30522
+MIN_SAFE_VOCAB_SIZE = 100
+FIXED_SAFE_LIMIT = 25000  # For very large vocabularies or specific model types
+LARGE_VOCAB_THRESHOLD = 100000
+MEDIUM_VOCAB_THRESHOLD = 50000
+SMALL_VOCAB_THRESHOLD = 10000
+LARGE_VOCAB_RATIO = 0.7
+MEDIUM_VOCAB_RATIO = 0.8
+SMALL_VOCAB_RATIO = 0.9
+
+
+class TemplateCodeGenerator(BaseCodeGenerator):
+    """Code generator using Jinja2 template"""
+    
+    def __init__(self, template_path: Optional[str] = None):
+        """
+        Args:
+            template_path: Path to Jinja2 template file (optional)
+        """
+        self.template_path = template_path
+        self._template = None
+    
+    def generate(
+        self,
+        model_dir: Path,
+        model_metadata: ModelMetadata,
+        output_dir: Path,
+    ) -> Path:
+        """
+        Generate run_model.py extraction script using template
+        
+        Args:
+            model_dir: Path to model directory
+            model_metadata: Model metadata extracted from configuration
+            output_dir: Output directory for generated script
+            
+        Returns:
+            Path to generated script file
+        """
+        try:
+            output_dir.mkdir(parents=True, exist_ok=True)
+            code = self._generate_code(model_dir, model_metadata)
+            
+            script_path = output_dir / "run_model.py"
+            with open(script_path, 'w', encoding='utf-8') as f:
+                f.write(code)
+            
+            return script_path
+        except Exception as e:
+            raise CodeGenError(f"Failed to generate code: {e}") from e
+    
+    def _generate_code(self, model_dir: Path, model_metadata: ModelMetadata) -> str:
+        """Generate complete extraction script code string"""
+        # Generate model loading code
+        load_code = self._generate_model_loader(model_dir, model_metadata)
+        
+        # Generate input construction code
+        input_code = self._generate_input_code(model_metadata)
+        
+        # Generate main code
+        code = f'''import torch
+try:
+    from transformers import AutoModel
+except ImportError:
+    raise ImportError("transformers is required. Install with: pip install transformers")
+
+import graph_net
+
+def main():
+    # Load model
+{self._indent(load_code, 4)}
+    
+    # Prepare inputs
+{self._indent(input_code, 4)}
+    
+    # Extract graph
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = model.to(device).eval()
+    
+    # Move inputs to same device as model
+    inputs = {{k: v.to(device) for k, v in inputs.items()}}
+    
+    wrapped = graph_net.torch.extract(name="{model_metadata.model_id}", dynamic=True)(model).eval()
+    
+    with torch.no_grad():
+        wrapped(**inputs)
+
+if __name__ == "__main__":
+    main()
+'''
+        return code
+    
+    def _generate_model_loader(self, model_dir: Path, model_metadata: ModelMetadata) -> str:
+        """Generate model loading code based on model type"""
+        model_path = str(model_dir).replace("\\", "/")
+        
+        if model_metadata.model_type in ['bert', 'gpt', 't5', 'roberta']:
+            return f'model = AutoModel.from_pretrained("{model_path}")'
+        elif model_metadata.model_type in ['resnet', 'vgg', 'densenet']:
+            return f'model = torchvision.models.{model_metadata.model_type}(pretrained=True)'
+        else:
+            # Generic loading
+            return f'model = AutoModel.from_pretrained("{model_path}")'
+    
+    def _generate_input_code(self, model_metadata: ModelMetadata) -> str:
+        """Generate input tensor construction code based on model metadata"""
+        lines = ['inputs = {}']
+        
+        for name, shape in model_metadata.input_shapes.items():
+            dtype = model_metadata.input_dtypes.get(name, "int64")
+            torch_dtype = self._get_torch_dtype(dtype)
+            shape_tuple = f"({', '.join(map(str, shape))})"
+            
+            if dtype == "int64":
+                if "input_ids" in name.lower():
+                    safe_vocab_size = self._calculate_safe_vocab_size(model_metadata)
+                    lines.append(f'inputs["{name}"] = torch.randint(0, {safe_vocab_size}, {shape_tuple}, dtype={torch_dtype})')
+                else:
+                    lines.append(f'inputs["{name}"] = torch.ones({shape_tuple}, dtype={torch_dtype})')
+            else:
+                lines.append(f'inputs["{name}"] = torch.randn({shape_tuple}, dtype={torch_dtype})')
+        
+        return "\n".join(lines)
+    
+    def _get_torch_dtype(self, dtype: str) -> str:
+        """Convert dtype string to torch dtype"""
+        if dtype == "int64":
+            return "torch.int64"
+        elif dtype == "float32":
+            return "torch.float32"
+        else:
+            return f"torch.{dtype}"
+    
+    def _calculate_safe_vocab_size(self, model_metadata: ModelMetadata) -> int:
+        """Calculate safe vocabulary size for input generation"""
+        if model_metadata.embedding_size:
+            return max(MIN_SAFE_VOCAB_SIZE, model_metadata.embedding_size - 1)
+        
+        vocab_size = model_metadata.vocab_size or DEFAULT_VOCAB_SIZE
+        model_type = (model_metadata.model_type or "").lower()
+        
+        # Model-type-specific limits
+        if self._is_large_vocab_model_type(model_type):
+            return FIXED_SAFE_LIMIT
+        
+        # Size-based strategy
+        if vocab_size > LARGE_VOCAB_THRESHOLD:
+            return FIXED_SAFE_LIMIT
+        elif vocab_size > MEDIUM_VOCAB_THRESHOLD:
+            return max(MIN_SAFE_VOCAB_SIZE, int(vocab_size * LARGE_VOCAB_RATIO))
+        elif vocab_size > SMALL_VOCAB_THRESHOLD:
+            return max(MIN_SAFE_VOCAB_SIZE, int(vocab_size * MEDIUM_VOCAB_RATIO))
+        else:
+            return max(MIN_SAFE_VOCAB_SIZE, int(vocab_size * SMALL_VOCAB_RATIO))
+    
+    def _is_large_vocab_model_type(self, model_type: str) -> bool:
+        """Check if model type typically has large vocabulary but small embedding"""
+        return "xlm-roberta" in model_type or "xlm_roberta" in model_type or "roberta" in model_type
+    
+    def _indent(self, text: str, spaces: int) -> str:
+        """Indent text by specified spaces"""
+        indent = " " * spaces
+        return "\n".join(indent + line for line in text.split("\n"))

graph_net/agent/graph_extractor/__init__.py

@@ -0,0 +1,6 @@
+"""Computation graph extraction modules"""
+
+from graph_net.agent.graph_extractor.base import BaseGraphExtractor
+from graph_net.agent.graph_extractor.subprocess_graph_extractor import SubprocessGraphExtractor
+
+__all__ = ["BaseGraphExtractor", "SubprocessGraphExtractor"]

graph_net/agent/graph_extractor/base.py

@@ -0,0 +1,25 @@
+"""Base graph extractor interface"""
+
+from abc import ABC, abstractmethod
+from pathlib import Path
+
+
+class BaseGraphExtractor(ABC):
+    """Base class for computation graph extractors"""
+    
+    @abstractmethod
+    def extract(self, code_path: Path, model_id: str) -> Path:
+        """
+        Execute script and extract computation graph
+        
+        Args:
+            code_path: Path to run_model.py script
+            model_id: Model ID for output directory naming
+            
+        Returns:
+            Path to extracted sample directory
+            
+        Raises:
+            ExtractionError: If extraction fails
+        """
+        pass