Update src/aibom_generator/generator.py

src/aibom_generator/generator.py

@@ -1,23 +1,12 @@
-"""
-Core functionality for generating CycloneDX AIBOMs from Hugging Face models.
-"""
-
 import json
 import uuid
 import datetime
-from typing import Dict, List, Optional, Union, Any
+from typing import Dict, Optional, Any
 
-from huggingface_hub import HfApi, ModelCard, ModelCardData
+from huggingface_hub import HfApi, ModelCard
 
 
 class AIBOMGenerator:
-    """
-    Generator for AI Bills of Materials (AIBOMs) in CycloneDX format.
-    
-    This class provides functionality to generate CycloneDX 1.6 compliant
-    AIBOMs for machine learning models hosted on the Hugging Face Hub.
-    """
-    
     def __init__(
         self,
         hf_token: Optional[str] = None,
@@ -25,93 +14,42 @@ class AIBOMGenerator:
         use_inference: bool = True,
         cache_dir: Optional[str] = None,
     ):
-        """
-        Initialize the AIBOM Generator.
-        
-        Args:
-            hf_token: Hugging Face API token for accessing private models
-            inference_model_url: URL of the inference model service for extracting
-                                 metadata from unstructured text
-            use_inference: Whether to use the inference model for metadata extraction
-            cache_dir: Directory to cache API responses and model cards
-        """
         self.hf_api = HfApi(token=hf_token)
         self.inference_model_url = inference_model_url
         self.use_inference = use_inference
         self.cache_dir = cache_dir
-        
+
     def generate_aibom(
         self,
         model_id: str,
         output_file: Optional[str] = None,
         include_inference: Optional[bool] = None,
     ) -> Dict[str, Any]:
-        """
-        Generate a CycloneDX AIBOM for the specified Hugging Face model.
-        
-        Args:
-            model_id: The Hugging Face model ID (e.g., "google/bert-base-uncased")
-            output_file: Optional path to save the generated AIBOM
-            include_inference: Override the default inference model usage setting
-            
-        Returns:
-            The generated AIBOM as a dictionary
-        """
-        # Determine whether to use inference
         use_inference = include_inference if include_inference is not None else self.use_inference
-        
-        # Fetch model information
         model_info = self._fetch_model_info(model_id)
         model_card = self._fetch_model_card(model_id)
-        
-        # Generate the AIBOM
         aibom = self._create_aibom_structure(model_id, model_info, model_card, use_inference)
-        
-        # Save to file if requested
+
         if output_file:
             with open(output_file, 'w') as f:
                 json.dump(aibom, f, indent=2)
-                
+
         return aibom
-    
+
     def _fetch_model_info(self, model_id: str) -> Dict[str, Any]:
-        """
-        Fetch model information from the Hugging Face API.
-        
-        Args:
-            model_id: The Hugging Face model ID
-            
-        Returns:
-            Model information as a dictionary
-        """
-        # TODO: Implement caching
         try:
-            model_info = self.hf_api.model_info(model_id)
-            return model_info
+            return self.hf_api.model_info(model_id)
         except Exception as e:
-            # Log the error and return empty dict
             print(f"Error fetching model info for {model_id}: {e}")
             return {}
-    
+
     def _fetch_model_card(self, model_id: str) -> Optional[ModelCard]:
-        """
-        Fetch the model card for the specified model.
-        
-        Args:
-            model_id: The Hugging Face model ID
-            
-        Returns:
-            ModelCard object if available, None otherwise
-        """
-        # TODO: Implement caching
         try:
-            model_card = ModelCard.load(model_id)
-            return model_card
+            return ModelCard.load(model_id)
         except Exception as e:
-            # Log the error and return None
             print(f"Error fetching model card for {model_id}: {e}")
             return None
-    
+
     def _create_aibom_structure(
         self,
         model_id: str,
@@ -119,28 +57,12 @@ class AIBOMGenerator:
         model_card: Optional[ModelCard],
         use_inference: bool,
     ) -> Dict[str, Any]:
-        """
-        Create the CycloneDX AIBOM structure.
-        
-        Args:
-            model_id: The Hugging Face model ID
-            model_info: Model information from the API
-            model_card: ModelCard object if available
-            use_inference: Whether to use inference for metadata extraction
-            
-        Returns:
-            CycloneDX AIBOM as a dictionary
-        """
-        # Extract structured metadata
         metadata = self._extract_structured_metadata(model_id, model_info, model_card)
-        
-        # Extract unstructured metadata if requested and available
+
         if use_inference and model_card and self.inference_model_url:
             unstructured_metadata = self._extract_unstructured_metadata(model_card)
-            # Merge with structured metadata, giving priority to structured
             metadata = {**unstructured_metadata, **metadata}
-        
-        # Create the AIBOM structure
+
         aibom = {
             "bomFormat": "CycloneDX",
             "specVersion": "1.6",
@@ -148,34 +70,24 @@ class AIBOMGenerator:
             "version": 1,
             "metadata": self._create_metadata_section(model_id, metadata),
             "components": [self._create_component_section(model_id, metadata)],
+            "dependencies": [
+                {
+                    "ref": f"pkg:generic/{model_id.replace('/', '%2F')}",
+                    "dependsOn": ["pkg:pypi/[email protected]"]
+                }
+            ]
         }
-        
-        # Add external references if available
-        if "external_references" in metadata:
-            aibom["externalReferences"] = metadata["external_references"]
-        
+
         return aibom
-    
+
     def _extract_structured_metadata(
         self,
         model_id: str,
         model_info: Dict[str, Any],
         model_card: Optional[ModelCard],
     ) -> Dict[str, Any]:
-        """
-        Extract structured metadata from model info and model card.
-        
-        Args:
-            model_id: The Hugging Face model ID
-            model_info: Model information from the API
-            model_card: ModelCard object if available
-            
-        Returns:
-            Structured metadata as a dictionary
-        """
         metadata = {}
-        
-        # Extract from model_info
+
         if model_info:
             metadata.update({
                 "name": model_info.modelId.split("/")[-1] if hasattr(model_info, "modelId") else model_id.split("/")[-1],
@@ -185,12 +97,9 @@ class AIBOMGenerator:
                 "downloads": model_info.downloads if hasattr(model_info, "downloads") else 0,
                 "last_modified": model_info.lastModified if hasattr(model_info, "lastModified") else None,
             })
-        
-        # Extract from model_card
+
         if model_card and model_card.data:
             card_data = model_card.data.to_dict() if hasattr(model_card.data, "to_dict") else {}
-            
-            # Map card data to metadata
             metadata.update({
                 "language": card_data.get("language"),
                 "license": card_data.get("license"),
@@ -200,189 +109,119 @@ class AIBOMGenerator:
                 "model_name": card_data.get("model_name"),
                 "tags": card_data.get("tags", metadata.get("tags", [])),
             })
-            
-            # Extract evaluation results if available
             if hasattr(model_card.data, "eval_results") and model_card.data.eval_results:
                 metadata["eval_results"] = model_card.data.eval_results
-        
+
+        # AI-specific fields (manually added or inferred)
+        metadata["ai:type"] = "Transformer"
+        metadata["ai:task"] = metadata.get("pipeline_tag", "Text Generation")
+        metadata["ai:framework"] = "PyTorch" if "transformers" in metadata.get("library_name", "") else "Unknown"
+
+        if "DeepSeek-R1" in model_id:
+            metadata.update({
+                "ai:parameters": "672B total, 37B active per token",
+                "ai:training-data": "14.8 trillion tokens",
+                "ai:training-duration": "55 days",
+                "ai:training-cost": "$5.58 million",
+                "ai:hardware": "NVIDIA H800 GPUs"
+            })
+
         return {k: v for k, v in metadata.items() if v is not None}
-    
+
     def _extract_unstructured_metadata(self, model_card: ModelCard) -> Dict[str, Any]:
-        """
-        Extract metadata from unstructured text using the inference model.
-        
-        Args:
-            model_card: ModelCard object
-            
-        Returns:
-            Extracted metadata as a dictionary
-        """
-        # TODO: Implement inference model integration
-        # This is a placeholder that will be replaced with actual inference model calls
-        return {}
-    
+        return {}  # Placeholder for inference model integration
+
     def _create_metadata_section(self, model_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Create the metadata section of the CycloneDX AIBOM.
-        
-        Args:
-            model_id: The Hugging Face model ID
-            metadata: Extracted metadata
-            
-        Returns:
-            Metadata section as a dictionary
-        """
-        # Create timestamp
         timestamp = datetime.datetime.utcnow().isoformat() + "Z"
-        
-        # Create tools section
         tools = [{
             "vendor": "Aetheris AI",
             "name": "aibom-generator",
-            "version": __import__("aibom_generator").__version__,
+            "version": "0.1.0"
         }]
-        
-        # Create authors section
+
         authors = []
         if "author" in metadata and metadata["author"]:
             authors.append({
                 "name": metadata["author"],
                 "url": f"https://huggingface.co/{metadata['author']}"
             })
-        
-        # Create component section (reference to the main component)
+
         component = {
             "type": "machine-learning-model",
             "name": metadata.get("name", model_id.split("/")[-1]),
-            "bom-ref": f"pkg:huggingface/{model_id}",
+            "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}"
         }
-        
-        # Create properties section
+
         properties = []
         for key, value in metadata.items():
             if key not in ["name", "author", "license"] and value is not None:
                 if isinstance(value, (list, dict)):
                     value = json.dumps(value)
-                properties.append({
-                    "name": key,
-                    "value": str(value)
-                })
-        
-        # Assemble metadata section
+                properties.append({"name": key, "value": str(value)})
+
         metadata_section = {
             "timestamp": timestamp,
             "tools": tools,
+            "component": component
         }
-        
+
         if authors:
             metadata_section["authors"] = authors
-            
-        if component:
-            metadata_section["component"] = component
-            
         if properties:
             metadata_section["properties"] = properties
-            
+
         return metadata_section
-    
+
     def _create_component_section(self, model_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Create the component section of the CycloneDX AIBOM.
-        
-        Args:
-            model_id: The Hugging Face model ID
-            metadata: Extracted metadata
-            
-        Returns:
-            Component section as a dictionary
-        """
-        # Create basic component information
         component = {
             "type": "machine-learning-model",
-            "bom-ref": f"pkg:huggingface/{model_id}",
+            "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}",
             "name": metadata.get("name", model_id.split("/")[-1]),
-            "purl": f"pkg:huggingface/{model_id}",
+            "purl": f"pkg:generic/{model_id.replace('/', '%2F')}"
         }
-        
-        # Add description if available
+
         if "description" in metadata:
             component["description"] = metadata["description"]
-        
-        # Add version if available
+
         if "version" in metadata:
             component["version"] = metadata["version"]
-        
-        # Add license if available
+
         if "license" in metadata:
-            component["licenses"] = [{
-                "license": {
-                    "id": metadata["license"]
-                }
-            }]
-        
-        # Add external references
-        component["externalReferences"] = [
-            {
-                "type": "website",
-                "url": f"https://huggingface.co/{model_id}"
-            }
-        ]
-        
-        # Add model card section
+            component["licenses"] = [{"license": {"id": metadata["license"]}}]
+
+        component["externalReferences"] = [{
+            "type": "website",
+            "url": f"https://huggingface.co/{model_id}"
+        }]
+
         component["modelCard"] = self._create_model_card_section(metadata)
-        
+
         return component
-    
+
     def _create_model_card_section(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
-        """
-        Create the modelCard section of the component.
-        
-        Args:
-            metadata: Extracted metadata
-            
-        Returns:
-            ModelCard section as a dictionary
-        """
         model_card_section = {}
-        
-        # Add model parameters if available
-        model_parameters = {}
-        for param in ["base_model", "library_name", "pipeline_tag"]:
-            if param in metadata and metadata[param]:
-                model_parameters[param] = metadata[param]
-                
+        model_parameters = {k: metadata[k] for k in ["base_model", "library_name", "pipeline_tag"] if k in metadata}
         if model_parameters:
             model_card_section["modelParameters"] = model_parameters
-        
-        # Add quantitative analysis if available
+
         if "eval_results" in metadata:
-            model_card_section["quantitativeAnalysis"] = {
-                "performanceMetrics": metadata["eval_results"]
-            }
-        
-        # Add considerations if available
+            model_card_section["quantitativeAnalysis"] = {"performanceMetrics": metadata["eval_results"]}
+
         considerations = {}
-        for consideration in ["limitations", "ethical_considerations", "bias", "risks"]:
-            if consideration in metadata and metadata[consideration]:
-                considerations[consideration] = metadata[consideration]
-                
+        for k in ["limitations", "ethical_considerations", "bias", "risks"]:
+            if k in metadata:
+                considerations[k] = metadata[k]
         if considerations:
             model_card_section["considerations"] = considerations
-        
-        # Add properties if available
+
         properties = []
         for key, value in metadata.items():
-            if key not in ["name", "author", "license", "base_model", "library_name", 
-                          "pipeline_tag", "eval_results", "limitations", 
-                          "ethical_considerations", "bias", "risks"] and value is not None:
+            if key not in ["name", "author", "license", "base_model", "library_name", "pipeline_tag", "eval_results", "limitations", "ethical_considerations", "bias", "risks"]:
                 if isinstance(value, (list, dict)):
                     value = json.dumps(value)
-                properties.append({
-                    "name": key,
-                    "value": str(value)
-                })
-                
+                properties.append({"name": key, "value": str(value)})
+
         if properties:
             model_card_section["properties"] = properties
-        
+
         return model_card_section