Update src/aibom_generator/generator.py

src/aibom_generator/generator.py

@@ -14,21 +14,27 @@ class AIBOMGenerator:
         inference_model_url: Optional[str] = None,
         use_inference: bool = True,
         cache_dir: Optional[str] = None,
+        use_best_practices: bool = True,  # Added parameter for industry-neutral scoring
     ):
         self.hf_api = HfApi(token=hf_token)
         self.inference_model_url = inference_model_url
         self.use_inference = use_inference
         self.cache_dir = cache_dir
         self.enhancement_report = None  # Store enhancement report as instance variable
+        self.use_best_practices = use_best_practices  # Store best practices flag
 
     def generate_aibom(
         self,
         model_id: str,
         output_file: Optional[str] = None,
         include_inference: Optional[bool] = None,
+        use_best_practices: Optional[bool] = None,  # Added parameter for industry-neutral scoring
     ) -> Dict[str, Any]:
         try:
             use_inference = include_inference if include_inference is not None else self.use_inference
+            # Use method parameter if provided, otherwise use instance variable
+            use_best_practices = use_best_practices if use_best_practices is not None else self.use_best_practices
+            
             model_info = self._fetch_model_info(model_id)
             model_card = self._fetch_model_card(model_id)
             
@@ -38,8 +44,8 @@ class AIBOMGenerator:
             # Create initial AIBOM with original metadata
             original_aibom = self._create_aibom_structure(model_id, original_metadata)
             
-            # Calculate initial score
-            original_score = calculate_completeness_score(original_aibom)
+            # Calculate initial score with industry-neutral approach if enabled
+            original_score = calculate_completeness_score(original_aibom, validate=True, use_best_practices=use_best_practices)
             
             # Final metadata starts with original metadata
             final_metadata = original_metadata.copy() if original_metadata else {}
@@ -69,8 +75,8 @@ class AIBOMGenerator:
             # Create final AIBOM with potentially enhanced metadata
             aibom = self._create_aibom_structure(model_id, final_metadata)
             
-            # Calculate final score
-            final_score = calculate_completeness_score(aibom)
+            # Calculate final score with industry-neutral approach if enabled
+            final_score = calculate_completeness_score(aibom, validate=True, use_best_practices=use_best_practices)
             
             # Add score and enhancement info to metadata properties
             if "metadata" in aibom and "properties" not in aibom["metadata"]:
@@ -83,6 +89,17 @@ class AIBOMGenerator:
                 aibom["metadata"]["properties"].append({"name": "aibom:quality-breakdown", "value": json.dumps(final_score["section_scores"])})
                 aibom["metadata"]["properties"].append({"name": "aibom:max-scores", "value": json.dumps(final_score["max_scores"])})
                 
+                # Add completeness profile information if available (from industry-neutral approach)
+                if use_best_practices and "completeness_profile" in final_score:
+                    aibom["metadata"]["properties"].append({
+                        "name": "aibom:completeness-profile", 
+                        "value": final_score["completeness_profile"]["name"]
+                    })
+                    aibom["metadata"]["properties"].append({
+                        "name": "aibom:completeness-description", 
+                        "value": final_score["completeness_profile"]["description"]
+                    })
+                
                 # Add AI enhancement information
                 if ai_enhanced:
                     aibom["metadata"]["properties"].append({"name": "aibom:ai-enhanced", "value": "true"})
@@ -181,6 +198,23 @@ class AIBOMGenerator:
             ]
         }
 
+        # Add critical fields for industry-neutral scoring (silently aligned with SPDX)
+        if metadata and "last_modified" in metadata:
+            aibom["buildTime"] = metadata["last_modified"]
+            # Add releaseTime as well if available, otherwise use last_modified
+            aibom["releaseTime"] = metadata.get("release_date", metadata["last_modified"])
+        
+        # Add downloadLocation if available
+        if metadata and "commit_url" in metadata:
+            # Add external reference for downloadLocation
+            if "externalReferences" not in aibom:
+                aibom["externalReferences"] = []
+            
+            aibom["externalReferences"].append({
+                "type": "distribution",
+                "url": f"https://huggingface.co/{model_id}"
+            })
+
         return aibom
 
     def _extract_structured_metadata(
@@ -227,6 +261,13 @@ class AIBOMGenerator:
         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"
+        
+        # Add fields for industry-neutral scoring (silently aligned with SPDX)
+        metadata["primaryPurpose"] = metadata.get("ai:task", "Text Generation")
+        metadata["suppliedBy"] = metadata.get("author", "Unknown")
+        
+        # Add typeOfModel field
+        metadata["typeOfModel"] = metadata.get("ai:type", "Transformer")
 
         return {k: v for k, v in metadata.items() if v is not None}
 
@@ -270,6 +311,8 @@ class AIBOMGenerator:
                         limitations_section = card_text.split("## Limitations")[1].split("##")[0].strip()
                         if limitations_section:
                             enhanced_metadata["limitations"] = limitations_section
+                            # Map to industry-neutral field (silently aligned with SPDX)
+                            enhanced_metadata["limitation"] = limitations_section
                 
                 # Extract ethical considerations if present
                 if "ethical_considerations" not in enhanced_metadata:
@@ -278,6 +321,8 @@ class AIBOMGenerator:
                             section = card_text.split(heading)[1].split("##")[0].strip()
                             if section:
                                 enhanced_metadata["ethical_considerations"] = section
+                                # Map to industry-neutral field (silently aligned with SPDX)
+                                enhanced_metadata["safetyRiskAssessment"] = section
                                 break
                 
                 # Extract risks if present
@@ -286,6 +331,24 @@ class AIBOMGenerator:
                         risks_section = card_text.split("## Risks")[1].split("##")[0].strip()
                         if risks_section:
                             enhanced_metadata["risks"] = risks_section
+                
+                # Extract energy consumption if present (for industry-neutral scoring)
+                if "energy" not in enhanced_metadata:
+                    for heading in ["## Energy", "## Energy Consumption", "## Environmental Impact"]:
+                        if heading in card_text:
+                            section = card_text.split(heading)[1].split("##")[0].strip()
+                            if section:
+                                enhanced_metadata["energyConsumption"] = section
+                                break
+                
+                # Extract hyperparameters if present (for industry-neutral scoring)
+                if "hyperparameters" not in enhanced_metadata:
+                    for heading in ["## Hyperparameters", "## Training Hyperparameters", "## Model Hyperparameters"]:
+                        if heading in card_text:
+                            section = card_text.split(heading)[1].split("##")[0].strip()
+                            if section:
+                                enhanced_metadata["hyperparameter"] = section
+                                break
             except Exception as e:
                 print(f"Error extracting unstructured metadata: {e}")
         
@@ -345,81 +408,127 @@ class AIBOMGenerator:
             
         component = {
             "type": "machine-learning-model",
-            "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}",
             "name": metadata.get("name", model_id.split("/")[-1]) if metadata else model_id.split("/")[-1],
+            "bom-ref": f"pkg:generic/{model_id.replace('/', '%2F')}",
             "purl": purl
         }
-
-        if metadata and "description" in metadata:
-            component["description"] = metadata["description"]
-
-        if metadata and "commit" in metadata:
-            component["version"] = metadata["commit"]
-
-        if metadata and "license" in metadata:
-            component["licenses"] = [{"license": {"id": metadata["license"]}}]
-
-        external_refs = [{
-            "type": "website",
-            "url": f"https://huggingface.co/{model_id}"
-        }]
-        if metadata and "commit_url" in metadata:
-            external_refs.append({
-                "type": "vcs",
-                "url": metadata["commit_url"]
-            })
-        component["externalReferences"] = external_refs
-
-        component["modelCard"] = self._create_model_card_section(metadata)
-
-        return component
-
-    def _create_model_card_section(self, metadata: Dict[str, Any]) -> Dict[str, Any]:
-        model_card_section = {}
         
-        if not metadata:
-            return model_card_section
+        # Add description if available
+        if metadata and "description" in metadata and metadata["description"]:
+            component["description"] = metadata["description"]
             
-        try:
-            # Create a CycloneDX-compliant modelParameters structure
-            # Instead of adding custom properties directly, use a standardized structure
-            model_parameters = {
-                "description": "Model parameters and configuration"
-            }
+        # Add license if available
+        if metadata and "license" in metadata and metadata["license"]:
+            license_id = metadata["license"]
+            component["licenses"] = [{
+                "license": {
+                    "id": license_id
+                }
+            }]
             
-            # Add properties array for custom parameters to ensure schema compliance
-            parameter_properties = []
+        # Add model card if available
+        model_card = {}
+        
+        # Add model parameters
+        model_parameters = {}
+        if metadata:
+            for key in ["ai:type", "ai:task", "ai:framework", "base_model", "library_name"]:
+                if key in metadata and metadata[key]:
+                    if "properties" not in model_parameters:
+                        model_parameters["properties"] = []
+                    model_parameters["properties"].append({
+                        "name": key,
+                        "value": metadata[key]
+                    })
+                    
+            # Add datasets if available
+            if "datasets" in metadata and metadata["datasets"]:
+                model_parameters["datasets"] = []
+                try:
+                    if isinstance(metadata["datasets"], list):
+                        for dataset in metadata["datasets"]:
+                            model_parameters["datasets"].append({
+                                "name": dataset
+                            })
+                    elif isinstance(metadata["datasets"], str):
+                        model_parameters["datasets"].append({
+                            "name": metadata["datasets"]
+                        })
+                except Exception as e:
+                    print(f"Error processing datasets: {e}")
+                    
+        if model_parameters:
+            model_card["modelParameters"] = model_parameters
             
-            for k in ["base_model", "library_name", "pipeline_tag"]:
-                if k in metadata and metadata[k] is not None:
-                    try:
-                        if isinstance(metadata[k], (list, dict)):
-                            value = json.dumps(metadata[k])
-                        else:
-                            value = str(metadata[k])
-                        parameter_properties.append({
-                            "name": k,
-                            "value": value
+        # Add quantitative analysis if available
+        if metadata and "eval_results" in metadata and metadata["eval_results"]:
+            try:
+                quantitative_analysis = {
+                    "performanceMetrics": []
+                }
+                
+                eval_results = metadata["eval_results"]
+                if isinstance(eval_results, dict):
+                    for metric, value in eval_results.items():
+                        quantitative_analysis["performanceMetrics"].append({
+                            "type": metric,
+                            "value": str(value)
                         })
-                    except Exception as e:
-                        print(f"Error processing model parameter {k}: {e}")
+                elif isinstance(eval_results, list):
+                    for result in eval_results:
+                        if isinstance(result, dict) and "metric" in result and "value" in result:
+                            quantitative_analysis["performanceMetrics"].append({
+                                "type": result["metric"],
+                                "value": str(result["value"])
+                            })
+                            
+                if quantitative_analysis["performanceMetrics"]:
+                    model_card["quantitativeAnalysis"] = quantitative_analysis
+            except Exception as e:
+                print(f"Error processing evaluation results: {e}")
+                
+        # Add considerations if available
+        considerations = {}
+        if metadata:
+            # Technical limitations
+            if "limitations" in metadata and metadata["limitations"]:
+                considerations["technicalLimitations"] = metadata["limitations"]
+                
+            # Ethical considerations
+            if "ethical_considerations" in metadata and metadata["ethical_considerations"]:
+                considerations["ethicalConsiderations"] = metadata["ethical_considerations"]
+                
+            # Risks
+            if "risks" in metadata and metadata["risks"]:
+                considerations["risks"] = metadata["risks"]
+                
+            # Environmental considerations (for industry-neutral scoring)
+            if "energyConsumption" in metadata and metadata["energyConsumption"]:
+                considerations["environmentalConsiderations"] = metadata["energyConsumption"]
+                
+        if considerations:
+            model_card["considerations"] = considerations
             
-            if parameter_properties:
-                model_parameters["properties"] = parameter_properties
-                model_card_section["modelParameters"] = model_parameters
-
-            # Safely extract evaluation results
-            if "eval_results" in metadata:
-                model_card_section["quantitativeAnalysis"] = {"performanceMetrics": metadata["eval_results"]}
-
-            # Safely extract considerations
-            considerations = {}
-            for k in ["limitations", "ethical_considerations", "bias", "risks"]:
-                if k in metadata:
-                    considerations[k] = metadata[k]
-            if considerations:
-                model_card_section["considerations"] = considerations
-        except Exception as e:
-            print(f"Error creating model card section: {e}")
+        if model_card:
+            component["modelCard"] = model_card
             
-        return model_card_section
+        # Add external references if available
+        external_references = []
+        
+        # Add model card URL
+        external_references.append({
+            "type": "documentation",
+            "url": f"https://huggingface.co/{model_id}"
+        })
+        
+        # Add commit URL if available
+        if metadata and "commit_url" in metadata and metadata["commit_url"]:
+            external_references.append({
+                "type": "vcs",
+                "url": metadata["commit_url"]
+            })
+            
+        if external_references:
+            component["externalReferences"] = external_references
+            
+        return component