diff --git a/README.md b/README.md index 8d66ea7..796b863 100644 --- a/README.md +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 16a1cca..9c9622a 100644 --- a/pyproject.toml +++ b/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 = [ diff --git a/src/llm_dna/api.py b/src/llm_dna/api.py index 858e9a0..c5a60e9 100644 --- a/src/llm_dna/api.py +++ b/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 diff --git a/src/llm_dna/cli.py b/src/llm_dna/cli.py index 10d203f..ef98247 100644 --- a/src/llm_dna/cli.py +++ b/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 diff --git a/src/llm_dna/core/extraction.py b/src/llm_dna/core/extraction.py index 4560814..183e406 100644 --- a/src/llm_dna/core/extraction.py +++ b/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 diff --git a/src/llm_dna/models/ModelLoader.py b/src/llm_dna/models/ModelLoader.py index f3c8b65..bde650d 100644 --- a/src/llm_dna/models/ModelLoader.py +++ b/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-",