大模型上下文窗口管理技巧:突破长度限制的艺术

前言

大模型的上下文窗口(Context Window)是指模型能够处理的最大输入长度。目前主流模型的上下文窗口从 4K 到 128K 不等,GPT-4 Turbo 甚至达到了 128K tokens。然而,随着应用场景的复杂化,我们经常需要处理超过这个限制的文本。

在实际项目中,我曾遇到需要让模型分析长篇技术文档、处理长对话历史、或者对整本书进行摘要的场景。这些需求都涉及上下文窗口的管理问题。今天分享一些我在实践中总结的技巧。

上下文窗口的基础知识

Token 的概念

Token 是模型处理文本的基本单位。对于英文,一个 token 大约等于 4 个字符或 0.75 个单词。对于中文,通常 1-2 个汉字等于一个 token。

import tiktoken

# 使用 tiktoken 计算 token 数
enc = tiktoken.get_encoding("cl100k_base")  # GPT-4 使用

text = "Hello, world!"
tokens = enc.encode(text)
print(f"Text: {text}")
print(f"Tokens: {tokens}")
print(f"Token count: {len(tokens)}")

上下文窗口的组成

在对话场景中,上下文窗口通常包含:

┌────────────────────────────────────────────────────────────┐
│                    上下文窗口                               │
├────────────────────────────────────────────────────────────┤
│ 系统提示 │ 用户历史 │ 助手历史 │ 当前查询 │ 可用长度 │        │
│ (System) │ (User)   │ (Assistant) │ (Query)  │ for Reply  │
└────────────────────────────────────────────────────────────┘

系统提示和对话历史会占用上下文空间,我们需要管理好这个空间,确保有足够的空间生成回复。

分块策略

固定窗口分块

最简单的方法是将文本按固定大小分割:

class FixedSizeChunker:
    def __init__(self, chunk_size: int, overlap: int):
        self.chunk_size = chunk_size
        self.overlap = overlap
    
    def chunk(self, text: str) -> List[str]:
        """固定大小分块"""
        words = text.split()
        chunks = []
        
        for i in range(0, len(words), self.chunk_size - self.overlap):
            chunk = " ".join(words[i:i + self.chunk_size])
            chunks.append(chunk)
        
        return chunks

语义分块

固定分块可能把语义相关的句子切断。语义分块尝试在句子边界处分割:

import re

class SemanticChunker:
    def __init__(self, max_tokens: int):
        self.max_tokens = max_tokens
    
    def chunk(self, text: str) -> List[str]:
        """按语义边界分块"""
        # 按段落分割
        paragraphs = text.split('\n')
        
        chunks = []
        current_chunk = []
        current_size = 0
        
        for para in paragraphs:
            para_tokens = self._count_tokens(para)
            
            if current_size + para_tokens > self.max_tokens:
                # 保存当前 chunk
                if current_chunk:
                    chunks.append('\n'.join(current_chunk))
                
                # 如果段落本身太大,按句子分割
                if para_tokens > self.max_tokens:
                    sub_chunks = self._split_long_paragraph(para)
                    chunks.extend(sub_chunks)
                    current_chunk = []
                    current_size = 0
                else:
                    current_chunk = [para]
                    current_size = para_tokens
            else:
                current_chunk.append(para)
                current_size += para_tokens
        
        # 添加最后一个 chunk
        if current_chunk:
            chunks.append('\n'.join(current_chunk))
        
        return chunks
    
    def _count_tokens(self, text: str) -> int:
        import tiktoken
        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))
    
    def _split_long_paragraph(self, para: str) -> List[str]:
        """分割长段落"""
        sentences = re.split(r'([。.!?])', para)
        chunks = []
        current = ""
        current_size = 0
        
        for i in range(0, len(sentences) - 1, 2):
            sentence = sentences[i] + sentences[i + 1]
            sentence_size = self._count_tokens(sentence)
            
            if current_size + sentence_size > self.max_tokens:
                if current:
                    chunks.append(current)
                current = sentence
                current_size = sentence_size
            else:
                current += sentence
                current_size += sentence_size
        
        if current:
            chunks.append(current)
        
        return chunks

递归分块

class RecursiveChunker:
    """使用多级分隔符的递归分块"""
    
    def __init__(self, separators: List[str], max_tokens: int):
        self.separators = separators
        self.max_tokens = max_tokens
    
    def chunk(self, text: str) -> List[str]:
        """递归分割直到满足大小要求"""
        return self._chunk_recursive(text, 0)
    
    def _chunk_recursive(self, text: str, separator_idx: int) -> List[str]:
        if separator_idx >= len(self.separators):
            # 最后手段:直接按大小截断
            return self._fixed_chunk(text)
        
        separator = self.separators[separator_idx]
        parts = text.split(separator)
        
        result = []
        current = ""
        
        for part in parts:
            test = current + separator + part if current else part
            test_size = self._count_tokens(test)
            
            if test_size <= self.max_tokens:
                current = test
            else:
                if current:
                    result.append(current)
                
                # 递归处理超出大小的部分
                if self._count_tokens(part) > self.max_tokens:
                    sub_chunks = self._chunk_recursive(part, separator_idx + 1)
                    result.extend(sub_chunks[:-1])
                    current = sub_chunks[-1] if sub_chunks else ""
                else:
                    current = part
        
        if current:
            result.append(current)
        
        return result
    
    def _fixed_chunk(self, text: str) -> List[str]:
        enc = tiktoken.get_encoding("cl100k_base")
        tokens = enc.encode(text)
        chunks = []
        
        for i in range(0, len(tokens), self.max_tokens):
            chunk_tokens = tokens[i:i + self.max_tokens]
            chunks.append(enc.decode(chunk_tokens))
        
        return chunks
    
    def _count_tokens(self, text: str) -> int:
        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))

对话历史管理

总结压缩

class ConversationManager:
    """对话历史管理器"""
    
    def __init__(
        self,
        llm,
        max_tokens: int = 3000,
        summary_tokens: int = 500
    ):
        self.llm = llm
        self.max_tokens = max_tokens
        self.summary_tokens = summary_tokens
        self.messages = []
    
    def add_message(self, role: str, content: str):
        """添加消息"""
        self.messages.append({"role": role, "content": content})
        self._maybe_summarize()
    
    def _maybe_summarize(self):
        """检查是否需要总结历史"""
        total_tokens = sum(
            self._count_tokens(m["content"]) 
            for m in self.messages
        )
        
        if total_tokens > self.max_tokens:
            self._summarize_old_messages()
    
    def _summarize_old_messages(self):
        """总结旧消息"""
        # 保留最近的消息
        recent_messages = self.messages[-4:]  # 最近 2 轮对话
        
        # 将早期消息合并为总结
        old_messages = self.messages[:-4]
        if not old_messages:
            return
        
        summary_prompt = f"""请总结以下对话的要点:

{chr(10).join([f"{m['role']}: {m['content']}" for m in old_messages])}

总结(简洁明了,保留关键信息):
"""
        
        summary = self.llm.generate(summary_prompt)
        
        # 用总结替换旧消息
        self.messages = [
            {"role": "system", "content": f"早期对话摘要:{summary}"}
        ] + recent_messages
    
    def get_context(self) -> List[dict]:
        """获取当前上下文"""
        return self.messages
    
    def _count_tokens(self, text: str) -> int:
        enc = tiktoken.get_encoding("cl100k_base")
        return len(enc.encode(text))

分层记忆

class HierarchicalMemory:
    """分层记忆系统"""
    
    def __init__(self, llm, token_limits: dict):
        self.llm = llm
        self.token_limits = token_limits
        
        # 三层记忆
        self.working_memory = []      # 当前对话
        self.episodic_memory = []     # 最近会话
        self.semantic_memory = ""    # 长期知识
    
    def add_working(self, role: str, content: str):
        """添加到工作记忆"""
        self.working_memory.append({"role": role, "content": content})
        self._condense_working()
    
    def _condense_working(self):
        """压缩工作记忆"""
        tokens = self._count_tokens(
            "\n".join([m["content"] for m in self.working_memory])
        )
        
        if tokens > self.token_limits["working"]:
            # 保留最近的,压缩更早的
            recent = self.working_memory[-2:]
            older = self.working_memory[:-2]
            
            if older:
                summary = self._summarize(older)
                self.episodic_memory.append({
                    "type": "episode",
                    "content": summary
                })
            
            self.working_memory = recent
    
    def get_context(self) -> str:
        """构建完整上下文"""
        parts = []
        
        # 语义记忆
        if self.semantic_memory:
            parts.append(f"【背景知识】\n{self.semantic_memory}")
        
        # 情景记忆
        if self.episodic_memory:
            episodes = "\n".join([
                f"- {ep['content']}" 
                for ep in self.episodic_memory[-3:]
            ])
            parts.append(f"【近期会话】\n{episodes}")
        
        # 工作记忆
        if self.working_memory:
            dialogue = "\n".join([
                f"{m['role']}: {m['content']}"
                for m in self.working_memory
            ])
            parts.append(f"【当前对话】\n{dialogue}")
        
        return "\n\n".join(parts)

长文本处理策略

Map-Reduce

处理长文本的经典方法:

class MapReduceProcessor:
    """Map-Reduce 文本处理"""
    
    def __init__(self, llm, chunk_size: int = 3000):
        self.llm = llm
        self.chunk_size = chunk_size
    
    def process(self, task: str, document: str) -> str:
        """Map-Reduce 处理"""
        # 1. Map:将文档分块,每块独立处理
        chunks = self._chunk(document)
        chunk_summaries = []
        
        for i, chunk in enumerate(chunks):
            summary = self._map_step(task, chunk, i + 1, len(chunks))
            chunk_summaries.append(summary)
        
        # 2. Reduce:合并所有结果
        final_result = self._reduce_step(task, chunk_summaries)
        
        return final_result
    
    def _chunk(self, text: str) -> List[str]:
        """分块"""
        enc = tiktoken.get_encoding("cl100k_base")
        tokens = enc.encode(text)
        chunks = []
        
        for i in range(0, len(tokens), self.chunk_size):
            chunk_tokens = tokens[i:i + self.chunk_size]
            chunks.append(enc.decode(chunk_tokens))
        
        return chunks
    
    def _map_step(self, task: str, chunk: str, idx: int, total: int) -> str:
        """Map 步骤:处理单个 chunk"""
        prompt = f"""这是关于"{task}"的任务。
文档的第 {idx}/{total} 部分内容如下:

{chunk}

请提取与任务相关的信息,用简洁的语言描述:
"""
        return self.llm.generate(prompt)
    
    def _reduce_step(self, task: str, summaries: List[str]) -> str:
        """Reduce 步骤:合并所有结果"""
        combined = "\n\n".join(
            f"第 {i+1} 部分:{s}" 
            for i, s in enumerate(summaries)
        )
        
        prompt = f"""任务:{task}

以下是文档各部分的处理结果:

{combined}

请综合以上信息,给出完整的回答:
"""
        return self.llm.generate(prompt)

Sliding Window Attention

模拟 Transformer 的滑动窗口机制:

class SlidingWindowProcessor:
    """滑动窗口处理"""
    
    def __init__(self, llm, window_size: int = 3000, overlap: int = 500):
        self.llm = llm
        self.window_size = window_size
        self.overlap = overlap
    
    def process(self, task: str, document: str) -> str:
        """滑动窗口处理"""
        enc = tiktoken.get_encoding("cl100k_base")
        tokens = enc.encode(document)
        
        results = []
        step = self.window_size - self.overlap
        
        for i in range(0, len(tokens), step):
            window_tokens = tokens[i:i + self.window_size]
            window_text = enc.decode(window_tokens)
            
            # 提取相关信息
            result = self._process_window(task, window_text)
            results.append(result)
            
            if i + self.window_size >= len(tokens):
                break
        
        # 合并结果
        return self._merge_results(task, results)
    
    def _process_window(self, task: str, window: str) -> str:
        prompt = f"""任务:{task}

内容窗口:

{window}

提取与任务相关的所有信息:
"""
        return self.llm.generate(prompt)

位置编码与外推

什么是外推

外推(Extrapolation)是指模型处理超过训练长度序列的能力。例如,如果模型在 2048 tokens 上训练,理论上就不能处理 4096 tokens 的输入。

RoPE 旋转位置编码

RoPE 是现代 LLM 广泛使用的位置编码方案:

def precompute_rope_params(head_dim, max_seq_len, theta=10000.0):
    """预计算 RoPE 参数"""
    inv_freq = 1.0 / (theta ** (torch.arange(0, head_dim, 2).float() / head_dim))
    t = torch.arange(max_seq_len)
    freqs = torch.einsum("i,j->ij", t, inv_freq)
    emb = torch.cat((freqs, freqs), dim=-1)
    return torch.cos(emb), torch.sin(emb)

def apply_rope(x, cos, sin):
    """应用 RoPE"""
    x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :]
    return torch.cat([
        x1 * cos - x2 * sin,
        x1 * sin + x2 * cos
    ], dim=-1)

位置插值

位置插值(Position Interpolation)是一种扩展上下文窗口的方法:

def interpolate_positions(freqs, scale_factor):
    """位置插值"""
    # 将位置索引缩放,使新位置映射到原始范围
    scaled_freqs = freqs / scale_factor
    return scaled_freqs

总结

上下文窗口管理是大模型应用中的重要课题。通过合理的分块策略、历史管理、压缩技术,我们可以有效突破模型的长度限制。

关键要点:

  1. 选择合适的分块策略(固定、语义、递归)
  2. 对话历史要及时压缩和总结
  3. 长文档使用 Map-Reduce 等分解策略
  4. 关注位置编码的外推能力

希望这些技巧对大家的项目有所帮助。

Logo

汇聚全球AI编程工具,助力开发者即刻编程。

更多推荐