Update src/aibom_generator/generator.py

src/aibom_generator/generator.py

@@ -96,6 +96,7 @@ class AIBOMGenerator:
                 "pipeline_tag": model_info.pipeline_tag if hasattr(model_info, "pipeline_tag") else None,
                 "downloads": model_info.downloads if hasattr(model_info, "downloads") else 0,
                 "last_modified": model_info.lastModified if hasattr(model_info, "lastModified") else None,
+                "commit": model_info.sha[:7] if hasattr(model_info, "sha") and model_info.sha else None,
             })
 
         if model_card and model_card.data:
@@ -112,7 +113,6 @@ class AIBOMGenerator:
             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"
@@ -129,7 +129,7 @@ class AIBOMGenerator:
         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 {}  # Placeholder for inference model integration
+        return {}
 
     def _create_metadata_section(self, model_id: str, metadata: Dict[str, Any]) -> Dict[str, Any]:
         timestamp = datetime.datetime.utcnow().isoformat() + "Z"
@@ -183,8 +183,8 @@ class AIBOMGenerator:
         if "description" in metadata:
             component["description"] = metadata["description"]
 
-        if "version" in metadata:
-            component["version"] = metadata["version"]
+        if "commit" in metadata:
+            component["version"] = metadata["commit"]
 
         if "license" in metadata:
             component["licenses"] = [{"license": {"id": metadata["license"]}}]