企业官网建设流程全解析

1. 生成一个长为 L 的干草堆文本（重复的无关文档） 2. 在深度 D% 的位置插入一条针信息（如 "小明最喜欢的水果是榴莲"） 3. 构造问题："小明最喜欢的水果是什么？" 4. 评估模型的回答是否正确 5. 改变 L 和 D，重复测试，形成评估矩阵

┌─────────────────────────────────────────────────────┐ │ NIAH Test Engine │ ├────────────┬────────────┬────────────┬───────────────┤ │ Generators │ Inserters │ Evaluators │ Analyzers │ │ │ │ │ │ │ • Needle │ • DeepPos │ • Exact │ • ScoreMap │ │ • Haystack│ • Random │ • Fuzzy │ • Heatmap │ │ • Query │ • Multi │ • LLM-J │ • Report │ └────────────┴────────────┴────────────┴───────────────┘

Config → TestCases → PromptBuilder → ModelRunner → Scorer → Report

from dataclasses import dataclass, field from typing import Optional, List import json @dataclass class NeedleSpec: """针的定义""" content: str # 针的内容，如 "小明最喜欢的水果是榴莲" question: str # 对应的问题 answer: str # 标准答案 needle_type: str = "fact" # fact / number / code / instruction @dataclass class NIAHTestCase: """单个测试用例""" context_length: int # 上下文总长度（字符数） insertion_depth: float # 插入深度（0.0 ~ 1.0） haystack_text: str # 干草堆文本 full_prompt: str # 完整提示词 needle_spec: NeedleSpec # 针的定义 model_response: Optional[str] = None score: Optional[float] = None metadata: dict = field(default_factory=dict)

import random from typing import List, Optional class HaystackGenerator: """干草堆文本生成器""" @staticmethod def _sample_documents(source_docs: List[str], target_chars: int, seed: Optional[int] = None) -> str: """ 从源文档库中采样，拼接到目标长度 模拟真实海量文档场景 """ if seed is not None: random.seed(seed) result = [] total = 0 while total < target_chars: doc = random.choice(source_docs) result.append(doc) total += len(doc) # 精确截断到目标长度 combined = "".join(result) return combined[:target_chars] @staticmethod def _generate_repetitive_text(template: str, target_chars: int, sep: str = "\n\n") -> str: """ 使用模板重复生成（低干扰方案） 适合快速验证测试 """ repeat_count = target_chars // (len(template) + len(sep)) return sep.join([template] * repeat_count)[:target_chars] @staticmethod def _generate_mixed_haystack(target_chars: int, num_topics: int = 20, seed: int = 42) -> str: """ 生成混合主题的干草堆（中干扰方案） 每个段落谈论不同主题 """ random.seed(seed) topics = [ "天气预报", "体育赛事", "科技新闻", "美食烹饪", "旅游攻略", "历史文化", "金融市场", "教育政策", "医疗健康", "环境保护", "交通出行", "建筑设计", "音乐鉴赏", "电影评论", "哲学思考", "农业技术", "航天探索", "海洋生物", "地壳运动", "考古发现" ] paragraphs = [] total = 0 while total < target_chars: topic = random.choice(topics) para = f"关于{topic}：这是{topic}相关的第{len(paragraphs)+1}段讨论。" para += f"在这个段落中，我们探讨{topic}的最新发展和重要发现。" para += f"研究表明，{topic}领域在过去一年取得了显著进展。" para += f"专家建议关注{topic}对日常生活的影响。\n\n" paragraphs.append(para) total += len(para) combined = "".join(paragraphs) return combined[:target_chars]

import uuid from typing import Tuple class NeedleGenerator: """生成不同类型的针""" # 预定义的事实针模板 FACT_TEMPLATES = [ "在{document_name}中，{person}最喜欢的{category}是{value}。", "根据{document_name}记录，{entity}的{attribute}是{value}。", "{document_name}第{section}节指出：{observation}。", ] CODE_TEMPLATES = [ "密码：{password}", "API_KEY = \"{api_key}\"", "secret = \"{secret_value}\"", ] @classmethod def generate_fact_needle(cls, doc_id: str = None) -> Tuple[str, str, str]: """ 生成事实型针，返回 (content, question, answer) """ if doc_id is None: doc_id = f"报告_{uuid.uuid4().hex[:8]}" persons = ["小明", "小红", "张三", "李四", "王五", "赵六"] categories = ["水果", "颜色", "运动", "城市", "书籍", "电影"] values = [ ("榴莲", "芒果", "荔枝", "草莓", "西瓜"), ("红色", "蓝色", "绿色", "紫色", "黄色"), ("篮球", "游泳", "跑步", "滑雪", "骑行"), ("巴黎", "东京", "大理", "冰岛", "京都"), ("三体", "百年孤独", "活着", "围城", "红楼梦"), ("星际穿越", "盗梦空间", "千与千寻", "让子弹飞", "肖申克的救赎"), ] person = random.choice(persons) cat_idx = random.randint(0, len(categories) - 1) category = categories[cat_idx] value = random.choice(values[cat_idx]) content = cls.FACT_TEMPLATES[0].format( document_name=doc_id, person=person, category=category, value=value ) question = f"在{doc_id}中，{person}最喜欢的{category}是什么？" answer = value return content, question, answer @classmethod def generate_number_needle(cls) -> Tuple[str, str, str]: """ 生成数字型针（精度更高，更适合精确评分） """ year = random.randint(2020, 2029) month = random.randint(1, 12) day = random.randint(1, 28) amount = random.randint(1000, 99999) content = f"交易记录：订单 #{random.randint(100000,999999)}，金额 {amount} 元，日期 {year}年{month}月{day}日。" question = f"订单 #{999999} 的金额是多少？" # 需要从上下文匹配 answer = str(amount) return content, question, answer

class NIAHBuilder: """将针插入干草堆并构建完整提示词""" @staticmethod def insert_needle(haystack: str, needle_content: str, depth: float) -> str: """ 在干草堆的指定深度位置插入针 depth: 0.0（开头）~ 1.0（末尾） """ if depth < 0 or depth > 1: raise ValueError(f"depth must be in [0, 1], got {depth}") # 计算插入点 insert_pos = int(len(haystack) * depth) # 插入针 result = haystack[:insert_pos] + needle_content + haystack[insert_pos:] return result @staticmethod def build_prompt(haystack_with_needle: str, question: str, instruction_template: Optional[str] = None) -> str: """ 构建完整的模型提示词 """ if instruction_template is None: instruction_template = ( "以下是一组文档。请仔细阅读所有内容，然后回答最后的问题。\n" "回答要准确、简洁，只给出答案即可，不要解释。\n\n" "--- 文档开始 ---\n" "{context}\n" "--- 文档结束 ---\n\n" "问题：{question}\n\n" "答案：" ) return instruction_template.format( context=haystack_with_needle, question=question )

from itertools import product class NIAHConfig: """测试配置与用例生成""" def __init__(self, context_lengths: List[int] = None, depths: List[float] = None, num_repeats: int = 3, needle_generator: str = "fact"): self.context_lengths = context_lengths or [ 1000, 2000, 4000, 8000, 16000, 32000, 64000, 128000 ] self.depths = depths or [0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0] self.num_repeats = num_repeats self.needle_generator = needle_generator self.seed = 42 def generate_test_cases(self) -> List[NIAHTestCase]: """生成所有测试用例""" cases = [] random.seed(self.seed) for ctx_len, depth in product(self.context_lengths, self.depths): for repeat in range(self.num_repeats): # 每次测试使用不同的随机针 needle_spec = self._create_needle() # 生成干草堆 haystack = HaystackGenerator._generate_mixed_haystack( ctx_len, seed=self.seed + repeat ) # 插入针 haystack_with_needle = NIAHBuilder.insert_needle( haystack, needle_spec.content, depth ) # 构建提示词 prompt = NIAHBuilder.build_prompt( haystack_with_needle, needle_spec.question ) case = NIAHTestCase( context_length=ctx_len, insertion_depth=depth, haystack_text=haystack_with_needle, full_prompt=prompt, needle_spec=needle_spec, metadata={"repeat": repeat, "seed": self.seed + repeat} ) cases.append(case) return cases def _create_needle(self) -> NeedleSpec: content, question, answer = NeedleGenerator.generate_fact_needle() return NeedleSpec( content=content, question=question, answer=answer )

import re from difflib import SequenceMatcher class Scorer: """评估模型回答的正确性""" @staticmethod def exact_match(response: str, answer: str) -> float: """精确匹配""" return 1.0 if response.strip() == answer.strip() else 0.0 @staticmethod def contains_match(response: str, answer: str) -> float: """包含匹配：答案是否出现在回复中""" return 1.0 if answer.strip() in response.strip() else 0.0 @staticmethod def fuzzy_match(response: str, answer: str, threshold: float = 0.8) -> float: """模糊匹配：基于相似度""" ratio = SequenceMatcher(None, response.strip(), answer.strip()).ratio() return ratio if ratio >= threshold else 0.0 @staticmethod def number_match(response: str, answer: str) -> float: """数字匹配：提取所有数字并比较""" resp_nums = re.findall(r'\d+', response) ans_nums = re.findall(r'\d+', answer) if not ans_nums: return 0.0 # 比较数字列表 matched = sum(1 for n in ans_nums if n in resp_nums) return matched / len(ans_nums) @classmethod def score(cls, response: str, answer: str, needle_type: str = "fact") -> float: """智能选择评分策略""" if not response: return 0.0 if needle_type == "number": return cls.number_match(response, answer) # 先试精确 if cls.exact_match(response, answer) == 1.0: return 1.0 # 再试包含 if cls.contains_match(response, answer) == 1.0: return 1.0 # 最后模糊 return cls.fuzzy_match(response, answer)

import time from concurrent.futures import ThreadPoolExecutor, as_completed class ModelRunner: """并发调用模型 API""" def __init__(self, api_key: str, base_url: str, model_name: str): self.api_key = api_key self.base_url = base_url.rstrip("/") self.model_name = model_name self.session = None # 实际使用 requests.Session() def _call_single(self, prompt: str, timeout: int = 60) -> str: """调用单个模型的完整实现""" import requests headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" } payload = { "model": self.model_name, "messages": [ {"role": "user", "content": prompt} ], "temperature": 0.0, # 确定性输出 "max_tokens": 50 # 只需简短答案 } try: resp = requests.post( f"{self.base_url}/v1/chat/completions", headers=headers, json=payload, timeout=timeout ) resp.raise_for_status() result = resp.json() return result["choices"][0]["message"]["content"].strip() except Exception as e: return f"[ERROR] {str(e)}" def evaluate_batch(self, test_cases: List[NIAHTestCase], max_workers: int = 5) -> List[NIAHTestCase]: """批量评估测试用例""" with ThreadPoolExecutor(max_workers=max_workers) as executor: futures = {} for case in test_cases: future = executor.submit( self._call_single, case.full_prompt ) futures[future] = case time.sleep(0.1) # 避免 API 限流 for future in as_completed(futures): case = futures[future] try: response = future.result() case.model_response = response case.score = Scorer.score( response, case.needle_spec.answer, case.needle_spec.needle_type ) except Exception as e: case.model_response = f"[EXCEPTION] {str(e)}" case.score = 0.0 return test_cases

import numpy as np from collections import defaultdict class ResultAnalyzer: """分析 NIAH 测试结果""" @staticmethod def build_score_matrix(results: List[NIAHTestCase]) -> dict: """构建 {length: {depth: avg_score}} 矩阵""" matrix = defaultdict(lambda: defaultdict(list)) for case in results: matrix[case.context_length][case.insertion_depth].append(case.score or 0.0) # 聚合（取平均值） aggregated = {} for length, depth_dict in matrix.items(): aggregated[length] = {} for depth, scores in depth_dict.items(): aggregated[length][depth] = np.mean(scores) return aggregated @staticmethod def compute_overall_score(matrix: dict) -> float: """计算综合评分""" all_scores = [] for length_dict in matrix.values(): all_scores.extend(length_dict.values()) return np.mean(all_scores) if all_scores else 0.0 @staticmethod def find_dead_zone(matrix: dict, threshold: float = 0.5) -> list: """ 找出"死亡区域"：模型表现低于阈值的 (长度, 深度) 区域 """ dead_zones = [] for length, depth_dict in matrix.items(): for depth, score in depth_dict.items(): if score < threshold: dead_zones.append({ "length": length, "depth": depth, "score": score }) return dead_zones

import matplotlib.pyplot as plt import matplotlib.colors as colors class NIAHVisualizer: """生成专业的 NIAH 热力图""" @staticmethod def plot_heatmap(matrix: dict, model_name: str = "Unknown", save_path: Optional[str] = None): """ 绘制 NIAH 测试热力图 X轴：插入深度（0%~100%） Y轴：上下文长度 颜色：得分（绿=好，红=差） """ lengths = sorted(matrix.keys()) depths = sorted(matrix[next(iter(matrix))].keys()) data = np.zeros((len(lengths), len(depths))) for i, l in enumerate(lengths): for j, d in enumerate(depths): data[i, j] = matrix[l].get(d, 0.0) fig, ax = plt.subplots(figsize=(12, 8)) # 使用红绿渐变（RdYlGn） cmap = plt.cm.RdYlGn norm = colors.Normalize(vmin=0, vmax=1) im = ax.imshow(data, cmap=cmap, norm=norm, aspect='auto') # 标签 ax.set_xlabel("插入深度 (%)", fontsize=12) ax.set_ylabel("上下文长度 (tokens)", fontsize=12) ax.set_title(f"Needle in a Haystack — {model_name}", fontsize=14, fontweight='bold') # 刻度 depth_labels = [f"{int(d*100)}%" for d in depths] ax.set_xticks(range(len(depths))) ax.set_xticklabels(depth_labels, rotation=45) length_labels = [format_length(l) for l in lengths] ax.set_yticks(range(len(lengths))) ax.set_yticklabels(length_labels) # 在格子中显示数值 for i in range(len(lengths)): for j in range(len(depths)): val = data[i, j] text_color = 'white' if val < 0.4 else 'black' ax.text(j, i, f"{val:.1f}", ha='center', va='center', color=text_color, fontsize=9, fontweight='bold') # 颜色条 cbar = plt.colorbar(im, ax=ax, fraction=0.046, pad=0.04) cbar.set_label("准确率", rotation=270, labelpad=15) plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') plt.close() else: plt.show() @staticmethod def plot_comparison(model_matrices: dict, save_path: Optional[str] = None): """ 多模型对比图 model_matrices: {model_name: matrix_dict} """ n_models = len(model_matrices) fig, axes = plt.subplots(1, n_models, figsize=(5*n_models, 6)) if n_models == 1: axes = [axes] cmap = plt.cm.RdYlGn norm = colors.Normalize(vmin=0, vmax=1) for ax, (name, matrix) in zip(axes, model_matrices.items()): lengths = sorted(matrix.keys()) depths = sorted(matrix[next(iter(matrix))].keys()) data = np.zeros((len(lengths), len(depths))) for i, l in enumerate(lengths): for j, d in enumerate(depths): data[i, j] = matrix[l].get(d, 0.0) ax.imshow(data, cmap=cmap, norm=norm, aspect='auto') ax.set_title(name, fontsize=11) ax.set_xlabel("深度") ax.set_ylabel("长度") # 刻度 depth_labels = [f"{int(d*100)}%" for d in depths] ax.set_xticks(range(len(depths))) ax.set_xticklabels(depth_labels, rotation=45, fontsize=8) length_labels = [format_length(l) for l in lengths] ax.set_yticks(range(len(lengths))) ax.set_yticklabels(length_labels, fontsize=8) plt.tight_layout() if save_path: plt.savefig(save_path, dpi=150, bbox_inches='tight') plt.close() else: plt.show() def format_length(chars: int) -> str: """将字符数格式化为可读的字符串""" if chars >= 10000: return f"{chars//1000}K" elif chars >= 1000: return f"{chars/1000:.1f}K" return str(chars)

import json from datetime import datetime class NIAHPipeline: """完整的 NIAH 测试流水线""" def __init__(self, config: NIAHConfig): self.config = config self.results = {} self.matrices = {} def run_single_model(self, model_name: str, api_key: str, base_url: str, output_dir: str = "./niah_results") -> dict: """ 对单个模型执行完整测试 """ os.makedirs(output_dir, exist_ok=True) print(f"[{datetime.now()}] 开始测试模型: {model_name}") print(f"[{datetime.now()}] 生成测试用例...") # 1. 生成测试用例 test_cases = self.config.generate_test_cases() print(f" → 共 {len(test_cases)} 个测试用例") # 2. 创建模型运行器 runner = ModelRunner(api_key, base_url, model_name) # 3. 批量评估 print(f"[{datetime.now()}] 开始评估...") results = runner.evaluate_batch(test_cases, max_workers=5) # 4. 分析结果 analyzer = ResultAnalyzer() matrix = analyzer.build_score_matrix(results) overall = analyzer.compute_overall_score(matrix) dead_zones = analyzer.find_dead_zone(matrix) print(f"\n[{datetime.now()}] 测试完成!") print(f" → 综合得分: {overall:.3f}") print(f" → 死亡区域数: {len(dead_zones)}") # 5. 保存原始结果 raw_path = os.path.join(output_dir, f"{model_name}_raw.json") self._save_raw(results, raw_path) # 6. 生成热力图 vis_path = os.path.join(output_dir, f"{model_name}_heatmap.png") NIAHVisualizer.plot_heatmap(matrix, model_name, save_path=vis_path) # 7. 生成报告 report_path = os.path.join(output_dir, f"{model_name}_report.json") report = { "model": model_name, "overall_score": overall, "dead_zones": dead_zones, "matrix": {str(k): {str(dk): dv for dk, dv in v.items()} for k, v in matrix.items()}, "config": { "lengths": self.config.context_lengths, "depths": self.config.depths, "repeats": self.config.num_repeats }, "timestamp": datetime.now().isoformat() } with open(report_path, "w", encoding="utf-8") as f: json.dump(report, f, ensure_ascii=False, indent=2) self.results[model_name] = results self.matrices[model_name] = matrix return report def compare_models(self, output_dir: str = "./niah_results"): """多模型对比""" if len(self.matrices) < 2: print("需要至少 2 个模型才能对比") return comp_path = os.path.join(output_dir, "comparison.png") NIAHVisualizer.plot_comparison(self.matrices, save_path=comp_path) # 生成对比报告 comparison = {} for name, matrix in self.matrices.items(): comparison[name] = ResultAnalyzer.compute_overall_score(matrix) print("\n模型排名：") for name, score in sorted(comparison.items(), key=lambda x: -x[1]): print(f" {name}: {score:.3f}") def _save_raw(self, results, path: str): """保存原始结果""" data = [] for case in results: data.append({ "context_length": case.context_length, "depth": case.insertion_depth, "needle": case.needle_spec.content, "question": case.needle_spec.question, "expected_answer": case.needle_spec.answer, "model_response": case.model_response, "score": case.score, "metadata": case.metadata }) with open(path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False, indent=2)

if __name__ == "__main__": import os # 配置 config = NIAHConfig( context_lengths=[1000, 4000, 8000, 16000, 32000, 64000], depths=[0.0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0], num_repeats=3 ) pipeline = NIAHPipeline(config) # 测试 DeepSeek report = pipeline.run_single_model( model_name="deepseek-chat", api_key=os.getenv("DEEPSEEK_API_KEY"), base_url="https://api.deepseek.com", output_dir="./niah_results" ) # 对比其他模型（如果需要） # pipeline.run_single_model("gpt-4o", gpt_key, "https://api.openai.com") # pipeline.compare_models()

class MultiNeedleBuilder: """多针插入策略""" @staticmethod def insert_multiple_needles(haystack: str, needles: List[Tuple[str, float]]) -> str: """ 在干草堆中插入多根针 needles: [(content, depth), ...] """ result = haystack # 按深度排序，从浅到深插入 sorted_needles = sorted(needles, key=lambda x: x[1]) # 倒序插入（避免后续插入影响前面位置） for content, depth in reversed(sorted_needles): insert_pos = int(len(result) * depth) result = result[:insert_pos] + content + result[insert_pos:] return result

class MultiHopNeedleGenerator: """多跳推理针""" @staticmethod def generate_two_hop_needle() -> Tuple[list, str, str]: """ 生成两跳推理测试 返回 (needle_list, question, answer) """ cities = { "北京": {"所在省": "河北省", "人口": 2154}, "上海": {"所在省": "江苏省", "人口": 2475}, "广州": {"所在省": "广东省", "人口": 1868}, "成都": {"所在省": "四川省", "人口": 2094}, } city_a, city_b = random.sample(list(cities.keys()), 2) needle1 = f"城市手册说：{city_a}的{list(cities[city_a].keys())[0]}是{list(cities[city_a].values())[0]}。" needle2 = f"另一份资料说：{city_b}的{list(cities[city_b].keys())[1]}是{list(cities[city_b].values())[1]}万人。" question = f"什么省是{cities[city_b]['所在省']}？{city_b}的人口是多少？" answer = f"{cities[city_b]['所在省']}，{cities[city_b]['人口']}万人" return [needle1, needle2], question, answer

class DistractorNeedleGenerator: """干扰针生成器""" @staticmethod def generate_with_distractor() -> Tuple[list, str, str]: """ 生成一个正确针 + 一个干扰针 """ true_value = random.randint(1000, 99999) distractor_value = true_value + random.choice([-1, 1]) * random.randint(10, 100) needles = [ f"正确记录：预算金额为 {true_value} 元。", f"（注意：初稿中预算曾误写作 {distractor_value} 元，最终以 {true_value} 元为准。）" ] question = "最终的预算金额是多少元？" answer = str(true_value) return needles, question, answer

class ReproducibleNIAH: """确保测试可复现""" @staticmethod def set_seed(seed: int = 42): """全局设置随机种子""" import random import numpy as np random.seed(seed) np.random.seed(seed) @staticmethod def save_test_config(config: NIAHConfig, path: str): """保存完整配置，便于复现""" import json config_dict = { "context_lengths": config.context_lengths, "depths": config.depths, "num_repeats": config.num_repeats, "needle_generator": config.needle_generator, "seed": config.seed, "generator_version": "v1.0", "timestamp": datetime.now().isoformat() } with open(path, "w", encoding="utf-8") as f: json.dump(config_dict, f, ensure_ascii=False, indent=2)

class AdaptiveNIAH: """自适应 NIAH：自动找到性能拐点""" def binary_search_limit(self, model_runner: ModelRunner, min_len: int = 1000, max_len: int = 200000, depth: float = 0.5, threshold: float = 0.8, step: int = 1000) -> int: """ 二分查找模型的有效上下文长度 """ lo, hi = min_len, max_len best = min_len while lo < hi: mid = (lo + hi) // 2 score = self._test_at_length(model_runner, mid, depth) if score >= threshold: best = mid lo = mid + step else: hi = mid - step return best

niah_test_engine/ ├── __init__.py ├── config.py # NIAHConfig, NIAHTestCase, NeedleSpec ├── generators.py # HaystackGenerator, NeedleGenerator ├── builder.py # NIAHBuilder ├── runner.py # ModelRunner ├── scorer.py # Scorer ├── analyzer.py # ResultAnalyzer ├── visualizer.py # NIAHVisualizer ├── pipeline.py # NIAHPipeline ├── advanced/ # 高级策略 │ ├── multi_needle.py │ ├── multi_hop.py │ └── distractor.py ├── results/ # 输出目录 └── requirements.txt

# requirements.txt numpy>=1.24.0 matplotlib>=3.7.0 requests>=2.31.0