Update src/aibom_generator/generator.py

src/aibom_generator/generator.py

@@ -4,6 +4,7 @@ import datetime
 from typing import Dict, Optional, Any
 
 from huggingface_hub import HfApi, ModelCard
+from utils import calculate_completeness_score
 
 
 class AIBOMGenerator:
@@ -118,23 +119,16 @@ class AIBOMGenerator:
         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}
+        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]:
         return {}
 
     def _create_metadata_section(self, model_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
-        from utils import calculate_completeness_score  # Assumes utils.py is in same module
-        completeness_score = calculate_completeness_score({"metadata": metadata})
+        aibom_stub = {"metadata": metadata}  # Build stub for scoring
+        score_report = calculate_completeness_score(aibom_stub)
+
         timestamp = datetime.datetime.utcnow().isoformat() + "Z"
         tools = [{
             "vendor": "Aetheris AI",
@@ -162,7 +156,10 @@ class AIBOMGenerator:
                     value = json.dumps(value)
                 properties.append({"name": key, "value": str(value)})
 
-                properties.append({"name": "aibom:quality-score", "value": str(completeness_score)})
+        # Add quality scoring results
+        properties.append({"name": "aibom:quality-score", "value": str(score_report["total_score"])})
+        properties.append({"name": "aibom:quality-breakdown", "value": json.dumps(score_report["section_scores"])})
+        properties.append({"name": "aibom:field-checklist", "value": json.dumps(score_report["field_checklist"])})
 
         metadata_section = {
             "timestamp": timestamp,