ai-agent-book 精选快照（<2MB 代码与文档，来自 github.com/bojieli/ai-agent-book）

chapter2/attention_visualization/config.py

@@ -0,0 +1,64 @@
+"""
+Configuration for Attention Visualization
+"""
+
+import os
+from pathlib import Path
+from dotenv import load_dotenv
+
+# Load environment variables
+load_dotenv()
+
+# Model Configuration
+MODEL_NAME = os.getenv("MODEL_NAME", "Qwen/Qwen3-0.6B")
+MODEL_PATH = os.getenv("MODEL_PATH", None)  # Optional local model path
+
+# Device Configuration
+DEVICE = os.getenv("DEVICE", "auto")  # auto, cuda, mps, or cpu
+
+# Attention Configuration
+ATTENTION_LAYER_INDEX = int(os.getenv("ATTENTION_LAYER_INDEX", -1))  # -1 for last layer
+TRACK_ALL_LAYERS = os.getenv("TRACK_ALL_LAYERS", "false").lower() == "true"
+
+# Generation Configuration
+DEFAULT_MAX_NEW_TOKENS = int(os.getenv("MAX_NEW_TOKENS", 100))
+DEFAULT_TEMPERATURE = float(os.getenv("TEMPERATURE", 0.7))
+DEFAULT_TOP_P = float(os.getenv("TOP_P", 0.9))
+DEFAULT_REPETITION_PENALTY = float(os.getenv("REPETITION_PENALTY", 1.1))
+
+# Visualization Configuration
+VIZ_OUTPUT_DIR = Path(os.getenv("VIZ_OUTPUT_DIR", "visualizations"))
+VIZ_FORMATS = os.getenv("VIZ_FORMATS", "heatmap,flow,summary").split(",")
+VIZ_COLORMAP = os.getenv("VIZ_COLORMAP", "viridis")
+VIZ_FIGSIZE = tuple(map(int, os.getenv("VIZ_FIGSIZE", "14,10").split(",")))
+VIZ_DPI = int(os.getenv("VIZ_DPI", 150))
+
+# Logging Configuration
+LOG_LEVEL = os.getenv("LOG_LEVEL", "INFO")
+LOG_FILE = Path(os.getenv("LOG_FILE", "attention_viz.log"))
+
+# Output Configuration  
+RESULTS_DIR = Path(os.getenv("RESULTS_DIR", "results"))
+RESULTS_DIR.mkdir(exist_ok=True)
+VIZ_OUTPUT_DIR.mkdir(exist_ok=True)
+
+# Interactive Mode Configuration
+INTERACTIVE_MODE = os.getenv("INTERACTIVE_MODE", "true").lower() == "true"
+AUTO_VISUALIZE = os.getenv("AUTO_VISUALIZE", "true").lower() == "true"
+
+# Demo Configuration
+DEMO_PROMPTS = [
+    "What is the capital of France?",
+    "Explain photosynthesis in simple terms",
+    "Write a haiku about artificial intelligence",
+    "List three benefits of exercise",
+    "What is 25 * 4 + 10?",
+]
+
+# System Prompts for Different Modes
+SYSTEM_PROMPTS = {
+    "default": "You are a helpful AI assistant.",
+    "technical": "You are a technical expert AI assistant. Provide detailed and accurate technical information.",
+    "creative": "You are a creative AI assistant. Be imaginative and original in your responses.",
+    "concise": "You are a concise AI assistant. Provide brief, clear answers without unnecessary elaboration.",
+}