Adjust/Improve pipeline to be compatible with more models

README.md

@@ -28,6 +28,27 @@ pip install llm-dna
 
 Use `llm-dna` for install/package naming, and `llm_dna` for Python imports.
 
+Optional extras are available for model families that need additional runtime dependencies:
+
+```bash
+# Apple Silicon / MLX-backed models
+pip install "llm-dna[apple]"
+
+# Quantized HuggingFace models (bitsandbytes, GPTQ, compressed-tensors, optimum)
+pip install "llm-dna[quantization]"
+
+# Architecture-specific model families such as Mamba or TIMM-backed models
+pip install "llm-dna[model_families]"
+
+# Everything above
+pip install "llm-dna[full]"
+```
+
+Extra guidance:
+- `apple`: required for MLX and `mlx-community/*` style model families on Apple Silicon.
+- `quantization`: required for many GPTQ, bitsandbytes, and compressed-tensors model families.
+- `model_families`: required for specific architectures whose modeling code depends on packages like `mamba-ssm` or `timm`.
+
 ## Quick Start
 
 ```python

pyproject.toml

@@ -32,11 +32,42 @@ dependencies = [
     "wonderwords>=2.2.0",
     "openai>=1.0.0",
     "tiktoken>=0.7.0",
+    "python-dotenv>=1.0.0",
 ]
 
 [project.optional-dependencies]
 model_scraping = ["requests>=2.31.0"]
 
+apple = [
+    "mlx>=0.10.0",
+    "mlx-lm>=0.10.0",
+]
+
+quantization = [
+    "bitsandbytes>=0.46.1",
+    "autoawq>=0.2.0",
+    "auto-gptq>=0.5.0",
+    "optimum>=1.16.0",
+    "compressed-tensors>=0.1.0",
+]
+
+model_families = [
+    "mamba-ssm>=1.0.0",
+    "timm>=0.9.0",
+]
+
+full = [
+    "mlx>=0.10.0",
+    "mlx-lm>=0.10.0",
+    "bitsandbytes>=0.46.1",
+    "autoawq>=0.2.0",
+    "auto-gptq>=0.5.0",
+    "optimum>=1.16.0",
+    "compressed-tensors>=0.1.0",
+    "mamba-ssm>=1.0.0",
+    "timm>=0.9.0",
+]
+
 vllm = ["vllm>=0.4.0"]
 
 dev = [

src/llm_dna/api.py

@@ -642,6 +642,31 @@ def calc_dna(config: DNAExtractionConfig) -> DNAExtractionResult:
         )
         vector = _validate_signature(signature)
 
+        if config.save:
+            cached_responses = _load_cached_responses(response_path, expected_count=len(probe_texts))
+            if cached_responses is None:
+                logging.info(
+                    "Generating and saving responses for '%s' to %s to align single-model caching with batch mode.",
+                    config.model_name,
+                    response_path,
+                )
+                responses = _generate_responses_for_model(
+                    model_name=config.model_name,
+                    config=config,
+                    model_meta=model_meta,
+                    probe_texts=probe_texts,
+                    device=resolved_device,
+                    resolved_token=resolved_token,
+                    incremental_save_path=response_path,
+                )
+                _save_response_cache(
+                    path=response_path,
+                    model_name=config.model_name,
+                    dataset=config.dataset,
+                    prompts=probe_texts,
+                    responses=responses,
+                )
+
     elapsed_seconds = time.time() - start_time
 
     output_path: Optional[Path] = None

src/llm_dna/cli.py

@@ -8,6 +8,8 @@
 from pathlib import Path
 from typing import Iterable, List, Optional
 
+from dotenv import load_dotenv
+
 
 def _load_models_from_file(path: Path) -> List[str]:
     """Load model names from a file, one per line."""
@@ -183,6 +185,10 @@ def main(argv: Optional[Iterable[str]] = None) -> int:
     """Main CLI entrypoint for DNA extraction."""
     from .api import DNAExtractionConfig, calc_dna, calc_dna_parallel
 
+    project_root = Path(__file__).resolve().parents[2]
+    load_dotenv(project_root / ".env", override=False)
+    load_dotenv(override=False)
+
     args = parse_arguments(argv)
 
     # Resolve model names

src/llm_dna/core/extraction.py

@@ -630,6 +630,13 @@ def main():
         logging.info(f"DNA signature saved to: {output_path}")
 
         # Save summary
+        # Create safe args dict without sensitive information
+        safe_args = vars(args).copy()
+        # Remove sensitive fields that should not be saved to output files
+        sensitive_fields = ['token', 'OPENROUTER_API_KEY', 'OPENAI_API_KEY']
+        for field in sensitive_fields:
+            safe_args.pop(field, None)
+
         summary = {
             "model_name": args.model_name,
             "dataset": args.dataset,
@@ -642,7 +649,7 @@ def main():
             "signature_stats": signature.get_statistics(),
             "metadata": signature.metadata.__dict__,
             "output_file": str(output_path),
-            "args": vars(args)
+            "args": safe_args
         }
 
         # Keep summary filename model-only as well

src/llm_dna/models/ModelLoader.py

@@ -3,6 +3,7 @@
 """
 
 import os
+import json
 from typing import Optional, Dict, Any, Union
 from pathlib import Path
 import logging
@@ -14,10 +15,34 @@
 
 class ModelLoader:
     """Factory class for loading different types of LLMs."""
+    _openrouter_model_ids: Optional[set[str]] = None
 
     def __init__(self, config_dict: Optional[Dict[str, Any]] = None):
         self.logger = logging.getLogger(__name__)
         self.config_dict = config_dict or {}
+
+    @classmethod
+    def _load_openrouter_model_ids(cls) -> set[str]:
+        if cls._openrouter_model_ids is not None:
+            return cls._openrouter_model_ids
+
+        model_ids: set[str] = set()
+        config_path = Path(__file__).resolve().parents[3] / "configs" / "openrouter_llm_list.jsonl"
+        try:
+            with config_path.open("r", encoding="utf-8") as handle:
+                for line in handle:
+                    line = line.strip()
+                    if not line:
+                        continue
+                    record = json.loads(line)
+                    model_id = str(record.get("model_id", "")).strip().lower()
+                    if model_id:
+                        model_ids.add(model_id)
+        except Exception:
+            model_ids = set()
+
+        cls._openrouter_model_ids = model_ids
+        return model_ids
 
     def load_model(
         self,
@@ -75,16 +100,20 @@ def _detect_model_type(self, model_path_or_name: str) -> str:
             "openrouter:",
             "anthropic/claude-",
             "deepseek/",
-            "openai/gpt-",
+            "openai/gpt-3",
+            "openai/gpt-4",
             "google/gemini-",
-            "z-ai/",
             "x-ai/grok-",
             "cohere/command",
             "perplexity/",
         ]
+
         if any(model_lower.startswith(prefix) for prefix in openrouter_prefixes):
             return "openrouter"
 
+        if model_lower in self._load_openrouter_model_ids():
+            return "openrouter"
+
         # Check for Google Gemini model names
         gemini_prefixes = [
             "gemini-",