diff --git a/src/models/extractor.py b/src/models/extractor.py index e47f4b4..8d23716 100644 --- a/src/models/extractor.py +++ b/src/models/extractor.py @@ -10,7 +10,11 @@ from .schemas import DataSource, ConfidenceLevel, ExtractionResult from .registry import get_field_registry_manager -from .model_file_extractors import ModelFileExtractor, default_extractors +from .model_file_extractors import ( + ModelFileExtractor, + default_extractors, + extract_distribution_metadata, +) from ..utils.license_utils import LICENSE_MAPPING from .config_parsing import parse_config as _parse_hparams_from_config @@ -205,6 +209,19 @@ def _registry_driven_extraction(self, model_id: str, model_info: Dict[str, Any], extraction_method="model_file_header", ) + # Format-agnostic distribution metadata (model_file_size, runtime_requirement) + distribution_metadata = self._extract_distribution_metadata(model_id) + if distribution_metadata: + for key, value in distribution_metadata.items(): + if value is not None and metadata.get(key) is None: + metadata[key] = value + self.extraction_results[key] = ExtractionResult( + value=value, + source=DataSource.REPOSITORY_FILES, + confidence=ConfidenceLevel.HIGH, + extraction_method="repository_file_listing", + ) + # Always extract commit SHA if available (vital for BOM versioning) if 'commit' not in metadata: commit_sha = getattr(model_info, 'sha', None) @@ -216,6 +233,14 @@ def _registry_driven_extraction(self, model_id: str, model_info: Dict[str, Any], return metadata + def _extract_distribution_metadata(self, model_id: str) -> Dict[str, Any]: + """Extract repo-level distribution attributes (size, runtime requirement).""" + try: + return extract_distribution_metadata(model_id, self.hf_api) + except Exception as e: + logger.warning(f"Distribution metadata extraction failed for {model_id}: {e}") + return {} + def _extract_model_file_metadata(self, model_id: str) -> Dict[str, Any]: for extractor in self.model_file_extractors: try: diff --git a/src/models/field_registry.json b/src/models/field_registry.json index 0300a8a..769ad38 100644 --- a/src/models/field_registry.json +++ b/src/models/field_registry.json @@ -1709,6 +1709,84 @@ "reference_urls": { "genai_aibom_taxonomy": "https://github.com/GenAI-Security-Project/cyclonedx-property-taxonomy" } + }, + "model_file_size": { + "tier": "supplementary", + "weight": 1.0, + "category": "component_model_card", + "description": "Total size in bytes of the model weight files in the repository", + "jsonpath": "$.components[0].modelCard.properties[?(@.name=='owasp:aibom:distribution:modelFileSize')].value", + "aibom_generation": { + "location": "$.components[0].properties", + "rule": "include_if_available", + "source_fields": [ + "model_file_size" + ], + "validation": "optional", + "data_type": "integer" + }, + "extraction": { + "methods": [ + "repository_files" + ], + "source_priority": [ + "repository_files" + ] + }, + "scoring": { + "points": 1.0, + "required_for_profiles": [ + "advanced" + ], + "category_contribution": 0.033 + }, + "validation_message": { + "missing": "Missing supplementary field: model_file_size - total weight file size helpful for deployment planning", + "recommendation": "Add the total size in bytes of the model weight files" + }, + "reference_urls": { + "note": "Custom OWASP extension property (owasp:aibom:distribution namespace); not part of the official genai:aibom taxonomy", + "cyclonedx_properties": "https://cyclonedx.org/docs/1.6/json/#components_items_properties" + } + }, + "runtime_requirement": { + "tier": "supplementary", + "weight": 1.0, + "category": "component_model_card", + "description": "Runtime needed to execute the model based on its weight format: 'python' (loadable from a Python ML stack such as transformers/PyTorch) or 'llm-stack' (requires an LLM inference stack such as llama.cpp/Ollama for GGUF/GGML weights)", + "jsonpath": "$.components[0].modelCard.properties[?(@.name=='owasp:aibom:distribution:runtimeRequirement')].value", + "aibom_generation": { + "location": "$.components[0].properties", + "rule": "include_if_available", + "source_fields": [ + "runtime_requirement" + ], + "validation": "optional", + "data_type": "string" + }, + "extraction": { + "methods": [ + "repository_files" + ], + "source_priority": [ + "repository_files" + ] + }, + "scoring": { + "points": 1.0, + "required_for_profiles": [ + "advanced" + ], + "category_contribution": 0.033 + }, + "validation_message": { + "missing": "Missing supplementary field: runtime_requirement - runtime/execution environment helpful for deployment", + "recommendation": "Add the runtime requirement ('python' or 'llm-stack') based on the model weight format" + }, + "reference_urls": { + "note": "Custom OWASP extension property (owasp:aibom:distribution namespace); not part of the official genai:aibom taxonomy", + "cyclonedx_properties": "https://cyclonedx.org/docs/1.6/json/#components_items_properties" + } } } } \ No newline at end of file diff --git a/src/models/model_file_extractors.py b/src/models/model_file_extractors.py index 2279cbf..8157c91 100644 --- a/src/models/model_file_extractors.py +++ b/src/models/model_file_extractors.py @@ -1,13 +1,33 @@ import logging -from typing import Protocol, Dict, List, Union, runtime_checkable +import os +import re +from typing import Optional, Protocol, Dict, List, Union, runtime_checkable -from huggingface_hub import list_repo_files +from huggingface_hub import HfApi, list_repo_files from .gguf_metadata import fetch_gguf_metadata_from_repo, map_to_metadata as gguf_map_to_metadata from .safetensors_metadata import fetch_safetensors_metadata, map_to_metadata as st_map_to_metadata logger = logging.getLogger(__name__) +# Weight-file formats loadable directly from a Python ML stack +# (transformers / PyTorch / TensorFlow / ONNX Runtime, etc.). +PYTHON_WEIGHT_EXTENSIONS = { + ".safetensors", ".bin", ".pt", ".pth", ".ckpt", + ".h5", ".msgpack", ".onnx", ".tflite", ".pb", +} +# Weight-file formats that require a dedicated LLM inference stack +# (llama.cpp / Ollama / LM Studio, etc.) rather than a Python import. +LLM_STACK_WEIGHT_EXTENSIONS = {".gguf", ".ggml"} + +# Taxonomy values for the runtime_requirement attribute. +RUNTIME_PYTHON = "python" +RUNTIME_LLM_STACK = "llm-stack" + +# Matches HuggingFace shard suffixes like "-00001-of-00002" or ".00001-of-00003" +# so shards of the same model/quant variant can be grouped together. +_SHARD_SUFFIX_RE = re.compile(r"[-.]?\d{3,5}-of-\d{3,5}") + @runtime_checkable class ModelFileExtractor(Protocol): @@ -61,3 +81,56 @@ def extract_metadata(self, model_id: str) -> Dict[str, Union[str, int, dict]]: def default_extractors() -> List[ModelFileExtractor]: return [SafetensorsFileExtractor(), GGUFFileExtractor()] + + +def extract_distribution_metadata( + model_id: str, hf_api: Optional[HfApi] = None +) -> Dict[str, Union[str, int]]: + """Extract format-agnostic distribution metadata for a model repo. + + Computes two attributes from the repository file listing (with sizes): + * ``model_file_size`` — total size in bytes of all model weight files + * ``runtime_requirement`` — ``python`` if the weights load from a Python ML + stack (e.g. safetensors/PyTorch), else ``llm-stack`` if the only weights + are GGUF/GGML (require llama.cpp/Ollama-style runtimes). + + Unlike the per-format ``ModelFileExtractor`` implementations, this runs for + every repo regardless of weight format. Returns an empty dict on failure. + """ + api = hf_api or HfApi() + try: + info = api.model_info(model_id, files_metadata=True) + except Exception as e: + logger.warning(f"Distribution metadata fetch failed for {model_id}: {e}") + return {} + + total_size = 0 + has_python_weights = False + has_llm_stack_weights = False + + for sibling in getattr(info, "siblings", None) or []: + filename = getattr(sibling, "rfilename", "") or "" + ext = os.path.splitext(filename.lower())[1] + size = getattr(sibling, "size", None) + + if ext in PYTHON_WEIGHT_EXTENSIONS: + has_python_weights = True + elif ext in LLM_STACK_WEIGHT_EXTENSIONS: + has_llm_stack_weights = True + else: + continue # not a weight file — ignore for size + runtime classification + + if isinstance(size, int) and size > 0: + total_size += size + + metadata: Dict[str, Union[str, int]] = {} + if total_size > 0: + metadata["model_file_size"] = total_size + # A repo that ships Python-loadable weights can run in Python even if it also + # ships GGUF; only GGUF/GGML-exclusive repos require an LLM stack. + if has_python_weights: + metadata["runtime_requirement"] = RUNTIME_PYTHON + elif has_llm_stack_weights: + metadata["runtime_requirement"] = RUNTIME_LLM_STACK + + return metadata diff --git a/src/models/service.py b/src/models/service.py index f3b676b..0e11edb 100644 --- a/src/models/service.py +++ b/src/models/service.py @@ -777,6 +777,17 @@ def _create_model_card_section(self, metadata: Dict[str, Any]) -> Dict[str, Any] props.append({"name": "genai:aibom:modelcard:quantizationFileType", "value": str(q_dict["file_type"])}) taxonomy_mapped_keys.append("quantization") + # Distribution properties — custom OWASP namespace. These are NOT part of the + # official genai:aibom taxonomy, which forbids unofficial names under genai:aibom. + distribution_mapping = { + "model_file_size": "modelFileSize", + "runtime_requirement": "runtimeRequirement", + } + for p_key, p_name in distribution_mapping.items(): + if metadata.get(p_key) is not None: + props.append({"name": f"owasp:aibom:distribution:{p_name}", "value": str(metadata[p_key])}) + taxonomy_mapped_keys.append(p_key) + # Basic Fields we've already mapped to structured homes mapped_fields = [ "primaryPurpose", "typeOfModel", "suppliedBy", "intendedUse", diff --git a/src/templates/result.html b/src/templates/result.html index 32a1c47..610d712 100644 --- a/src/templates/result.html +++ b/src/templates/result.html @@ -173,11 +173,38 @@

📊  Model Card

aibom.components[0].modelCard.modelParameters.task else 'Not specified' }} {% endif %} + {# Pull out the distribution (deployment/hardware) properties for a friendly render #} + {% set dist = namespace(file_size=None, runtime=None) %} {% if aibom.components[0].modelCard.properties %} + {% for prop in aibom.components[0].modelCard.properties %} + {% if prop.name == 'owasp:aibom:distribution:modelFileSize' %}{% set dist.file_size = prop.value %}{% endif %} + {% if prop.name == 'owasp:aibom:distribution:runtimeRequirement' %}{% set dist.runtime = prop.value %}{% endif %} + {% endfor %} + {% endif %} + {% if dist.file_size or dist.runtime %} +
+ Deployment & Hardware: + + {% if dist.file_size %} + Model Size: {{ dist.file_size | int | filesizeformat }} + {% endif %} + {% if dist.runtime == 'python' %} + Runtime: Python ML stack (transformers/PyTorch) + Hardware: CPU-friendly, no extra runtime stack + {% elif dist.runtime == 'llm-stack' %} + Runtime: LLM runtime (llama.cpp / Ollama / LM Studio) + Hardware: needs an inference stack + more RAM/VRAM + {% endif %} + +
+ {% endif %} + {# Any remaining modelCard properties that are not deployment ones #} + {% set other_props = aibom.components[0].modelCard.properties | rejectattr('name', 'in', ['owasp:aibom:distribution:modelFileSize', 'owasp:aibom:distribution:runtimeRequirement']) | list if aibom.components[0].modelCard.properties else [] %} + {% if other_props %}
Additional Properties: - {% for prop in aibom.components[0].modelCard.properties %} + {% for prop in other_props %} {{ prop.name }}: {{ prop.value }} {% endfor %} diff --git a/src/utils/summarizer.py b/src/utils/summarizer.py index c4f2348..1cb441a 100644 --- a/src/utils/summarizer.py +++ b/src/utils/summarizer.py @@ -1,4 +1,5 @@ import logging +import os import re from typing import Optional, List, Dict, Any @@ -11,7 +12,7 @@ class LocalSummarizer: """ _tokenizer = None _model = None - _model_name = "google/flan-t5-small" + _model_name = "facebook/bart-large-cnn" @classmethod def _load_model(cls): @@ -100,6 +101,12 @@ def _score_candidate(text: str) -> float: score -= 50.0 if "install" in text_lower or "how to run" in text_lower or "pip install" in text_lower or "read our guide" in text_lower: score -= 30.0 + + # Penalize quantization / repackaging marketing (e.g. unsloth/TheBloke GGUF + # repos lead with quant claims instead of the base model's actual purpose). + if ("outperforms other leading" in text_lower or "dynamic 2.0" in text_lower + or "quants" in text_lower or "quantized version" in text_lower): + score -= 40.0 # Penalize table/code-heavy paragraphs and bullet points if text.count('|') > 5 or text.count('```') >= 1 or text.count('\n- ') > 2 or text.count('\n* ') > 2: @@ -119,10 +126,19 @@ def _clean_text(text: str) -> str: except Exception: pass + # Remove markdown heading markers (keep the heading text) + text = re.sub(r'(?m)^\s{0,3}#{1,6}\s*', '', text) # Remove markdown images text = re.sub(r'!\[.*?\]\([^)]+\)', '', text) # Convert links to just text text = re.sub(r'\[([^\]]+)\]\([^)]+\)', r'\1', text) + # Remove markdown emphasis markers (bold/italic), keeping the inner text + text = re.sub(r'\*\*([^*]+)\*\*', r'\1', text) + text = re.sub(r'__([^_]+)__', r'\1', text) + text = re.sub(r'\*([^*]+)\*', r'\1', text) + text = re.sub(r'(? str: return text + @staticmethod + def _truncate(text: str, max_output_chars: int) -> str: + text = text.strip() + if len(text) > max_output_chars: + return text[:max_output_chars - 3].rstrip() + "..." + return text + + @staticmethod + def _strip_prompt_leak(text: str) -> str: + """Defensively remove any of our own instruction phrasing that a + non-instruction summarizer (BART) may echo back into its output.""" + patterns = [ + r'in one sentence,?\s*explain what this ai model is designed to do[^.]*\.?', + r'summarize the main purpose of this ai model[^.]*\.?', + r'based on this description:?', + ] + for pat in patterns: + text = re.sub(pat, '', text, flags=re.IGNORECASE) + return re.sub(r'\s+', ' ', text).strip() + + @staticmethod + def _use_ollama() -> bool: + """Use the instruction-tuned LLM backend when AIBOM_SUMMARIZER_BACKEND=ollama.""" + return os.environ.get("AIBOM_SUMMARIZER_BACKEND", "bart").strip().lower() == "ollama" + @classmethod - def _generate(cls, prompt: str, max_output_chars: int) -> Optional[str]: + def _generate_ollama(cls, source: str, max_output_chars: int) -> Optional[str]: + """Generate an objective purpose statement with an instruction-tuned LLM via Ollama. + + Unlike BART, this model *follows* instructions, so we ask it directly for an + objective, marketing-free purpose statement. Configured via env vars: + AIBOM_OLLAMA_HOST (default http://localhost:11434) + AIBOM_OLLAMA_MODEL (default qwen3:4b) + Returns None on any failure so summarize() can fall back to BART. + """ + import json + import urllib.request + + host = os.environ.get("AIBOM_OLLAMA_HOST", "http://localhost:11434").rstrip("/") + model = os.environ.get("AIBOM_OLLAMA_MODEL", "qwen3:4b") + + prompt = ( + "You are writing a factual purpose statement for an AI Bill of Materials (AIBOM).\n" + "From the model card below, state ONLY what the model does and its intended task.\n" + "Rules: be objective, use no marketing language, invent nothing, copy names and " + "numbers exactly, and respond with one or two sentences only. If the purpose is " + "not stated in the card, reply exactly: No description available.\n\n" + f"MODEL CARD:\n{source[:6000]}" + ) + payload = json.dumps({ + "model": model, + "prompt": prompt, + "stream": False, + "think": False, # disable chain-of-thought on reasoning models (e.g. qwen3) + "options": {"temperature": 0.0, "num_predict": 220}, + }).encode("utf-8") + + try: + req = urllib.request.Request( + f"{host}/api/generate", + data=payload, + headers={"Content-Type": "application/json"}, + ) + with urllib.request.urlopen(req, timeout=120) as resp: + data = json.loads(resp.read().decode("utf-8")) + except Exception as e: + logger.warning(f"⚠️ Ollama generation failed ({model} @ {host}): {e}") + return None + + summary = (data.get("response") or "").strip() + # Strip any leaked reasoning blocks from thinking models. + summary = re.sub(r'.*?', '', summary, flags=re.DOTALL).strip() + summary = cls._strip_prompt_leak(summary) + + if not summary or summary.strip().lower().rstrip('.') == "no description available": + return None + return cls._truncate(summary, max_output_chars) + + @classmethod + def _generate(cls, source: str, max_output_chars: int) -> Optional[str]: + """Summarize source text. + + NOTE: facebook/bart-large-cnn is a *summarization* model, not an + instruction-following LLM. The ``source`` passed here MUST be the + document text only — never an instruction/prompt. BART cannot tell an + instruction apart from the document and will copy it into the output + (the prompt-leak bug) and hallucinate generic "AI model designed to..." + summaries when primed with instruction words. + """ if cls._model is None: cls._load_model() if not cls._model or not cls._tokenizer: return None - + try: - inputs = cls._tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True) + inputs = cls._tokenizer(source, return_tensors="pt", max_length=512, truncation=True) generate_kwargs = { - "max_length": 128, # Increased by ~30% from 64 - "min_length": 15, # Avoid single word outputs + "max_length": 160, + "min_length": 20, # Avoid single word / fragment outputs "do_sample": False, "num_beams": 4, "early_stopping": True, - "repetition_penalty": 2.0 + "repetition_penalty": 2.0, + "no_repeat_ngram_size": 3, # Stop the repetitive looping seen on short inputs } summary_ids = cls._model.generate(inputs["input_ids"], **generate_kwargs) summary = cls._tokenizer.decode(summary_ids[0], skip_special_tokens=True) - + summary = summary.strip() - + # Remove "Output:" prefix if present if summary.lower().startswith("output:"): summary = re.sub(r'^Output:\s*', '', summary, flags=re.IGNORECASE) - - if len(summary) > max_output_chars: - return summary[:max_output_chars-3] + "..." - return summary + + summary = cls._strip_prompt_leak(summary) + + return cls._truncate(summary, max_output_chars) except Exception as e: logger.warning(f"⚠️ Generation failed: {e}") return None @@ -201,6 +305,10 @@ def _is_valid_summary(summary: str, model_id: str) -> bool: # Check for instruction-like text if summary_lower.startswith("to install") or summary_lower.startswith("how to") or "pip install" in summary_lower: return False + + # Reject leaked prompt/instruction phrasing (the prompt-leak bug) + if "explain what this ai model" in summary_lower or "summarize the main purpose" in summary_lower: + return False # Refuse literally copying bullet points (e.g. from table) if "- type:" in summary_lower or "number of parameters:" in summary_lower: @@ -209,7 +317,7 @@ def _is_valid_summary(summary: str, model_id: str) -> bool: return True @classmethod - def summarize(cls, text: str, max_output_chars: int = 332, model_id: str = "") -> Optional[str]: + def summarize(cls, text: str, max_output_chars: int = 1024, model_id: str = "") -> Optional[str]: """ Robustly extract and summarize model description. """ @@ -235,32 +343,39 @@ def summarize(cls, text: str, max_output_chars: int = 332, model_id: str = "") - if not cleaned_text.strip(): return None - - # Extract just the first few sentences of the cleaned text to avoid confusing the small model - # with training details that usually appear at the end of the paragraph. + sentences = re.split(r'(?<=[.!?])\s+', cleaned_text) - short_text = " ".join(sentences[:3]) - - # 5 & 6 & 7. Summarize, Validate, Retry, Fallback - prompt1 = f"In one sentence, explain what this AI model is designed to do based on this description:\n\n{short_text}" - - summary = cls._generate(prompt1, max_output_chars) - + + # Preferred backend: instruction-tuned LLM (Ollama). It follows instructions, + # so it produces an objective purpose statement directly. Falls through to BART + # if disabled, unavailable, or it returns an invalid summary. + if cls._use_ollama(): + summary = cls._generate_ollama(cleaned_text, max_output_chars) + if summary and cls._is_valid_summary(summary, model_id): + return summary + logger.info("⚠️ Ollama backend unavailable/invalid; falling back to BART.") + + # Use the first few sentences of the cleaned text as the source document. + # Trailing sentences are usually training/benchmark details that dilute the + # summary of the model's purpose. + source = " ".join(sentences[:5]).strip() + + # BART is a summarization model and hallucinates badly on very short inputs + # (looping repetition, invented "a new AI model designed to..." text). When we + # don't have enough source material to summarize, return the clean extract + # directly instead of letting the model invent a purpose. + MIN_CHARS_FOR_ABSTRACTION = 280 + if len(source) < MIN_CHARS_FOR_ABSTRACTION: + return cls._truncate(" ".join(sentences[:2]), max_output_chars) + + # Summarize the document text ONLY. Do not prepend an instruction/prompt: + # BART cannot follow instructions and would echo the prompt into the output + # (the prompt-leak bug) and produce generic hallucinations. + summary = cls._generate(source, max_output_chars) + if summary and cls._is_valid_summary(summary, model_id): return summary - - # Retry with stricter prompt - logger.info("⚠️ First summary invalid, retrying with stricter prompt.") - prompt2 = f"Summarize the main purpose of this AI model in one complete sentence:\n\n{cleaned_text}" - summary2 = cls._generate(prompt2, max_output_chars) - - if summary2 and cls._is_valid_summary(summary2, model_id): - return summary2 - - # Fallback to cleaned text (first 1-2 sentences) - logger.info("⚠️ Both LLM summaries invalid, falling back to cleaned extracted text.") - sentences = re.split(r'(?<=[.!?])\s+', cleaned_text) - fallback_summary = " ".join(sentences[:2]) - if len(fallback_summary) > max_output_chars: - return fallback_summary[:max_output_chars-3] + "..." - return fallback_summary + + # Fallback to the cleaned extracted text (first 1-2 sentences). + logger.info("⚠️ Summary invalid, falling back to cleaned extracted text.") + return cls._truncate(" ".join(sentences[:2]), max_output_chars)