14：1M+ Token长上下文在真实业务中的落地价值

class SparseAttention(nn.Module):
    def __init__(self, hidden_size, num_heads, top_k=64):
        super().__init__()
        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.top_k = top_k
        self.head_dim = hidden_size // num_heads
        
        self.q_proj = nn.Linear(hidden_size, hidden_size)
        self.k_proj = nn.Linear(hidden_size, hidden_size)
        self.v_proj = nn.Linear(hidden_size, hidden_size)
        self.out_proj = nn.Linear(hidden_size, hidden_size)
    
    def forward(self, x):
        batch_size, seq_len, hidden_size = x.shape
        
        # 计算查询、键、值
        q = self.q_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        k = self.k_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        v = self.v_proj(x).view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
        
        # 计算注意力分数
        attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)
        
        # 只保留top-k个注意力分数
        top_k_scores, top_k_indices = torch.topk(attn_scores, self.top_k, dim=-1)
        
        # 创建稀疏注意力矩阵
        sparse_attn = torch.zeros_like(attn_scores)
        sparse_attn.scatter_(-1, top_k_indices, top_k_scores)
        
        # 应用softmax
        sparse_attn = F.softmax(sparse_attn, dim=-1)
        
        # 计算加权和
        output = torch.matmul(sparse_attn, v)
        output = output.transpose(1, 2).contiguous().view(batch_size, seq_len, hidden_size)
        output = self.out_proj(output)
        
        return output

class LongContextManager:
    def __init__(self, max_context_length=1000000):
        self.max_context_length = max_context_length
        self.context = []
        self.token_count = 0
    
    def add_text(self, text, tokenizer):
        # 计算文本的token数量
        tokens = tokenizer.encode(text)
        token_count = len(tokens)
        
        # 如果添加后超过最大长度，移除最早的内容
        while self.token_count + token_count > self.max_context_length:
            removed_text = self.context.pop(0)
            removed_tokens = tokenizer.encode(removed_text)
            self.token_count -= len(removed_tokens)
        
        # 添加新文本
        self.context.append(text)
        self.token_count += token_count
        
        return self.get_full_context()
    
    def get_full_context(self):
        return " ".join(self.context)
    
    def get_relevant_context(self, query, tokenizer, max_relevant_length=100000):
        # 计算查询与每个文本片段的相关性
        query_tokens = tokenizer.encode(query)
        relevance_scores = []
        
        for text in self.context:
            text_tokens = tokenizer.encode(text)
            # 简单的相关性计算（实际应用中可以使用更复杂的方法）
            score = self.calculate_relevance(query_tokens, text_tokens)
            relevance_scores.append(score)
        
        # 按相关性排序
        sorted_indices = sorted(range(len(relevance_scores)), key=lambda i: relevance_scores[i], reverse=True)
        
        # 收集最相关的文本，直到达到最大长度
        relevant_context = []
        current_length = 0
        
        for i in sorted_indices:
            text = self.context[i]
            text_tokens = tokenizer.encode(text)
            text_length = len(text_tokens)
            
            if current_length + text_length <= max_relevant_length:
                relevant_context.append(text)
                current_length += text_length
            else:
                # 截断文本以适应长度
                remaining_length = max_relevant_length - current_length
                truncated_text = tokenizer.decode(text_tokens[:remaining_length])
                relevant_context.append(truncated_text)
                break
        
        return " ".join(relevant_context)
    
    def calculate_relevance(self, query_tokens, text_tokens):
        # 简单的重叠词计算
        query_set = set(query_tokens)
        text_set = set(text_tokens)
        intersection = query_set.intersection(text_set)
        return len(intersection) / len(query_set) if query_set else 0

class LongContextAnalyzer:
    def __init__(self, model_name="gemini-1.5-pro", max_context_length=1000000):
        self.model = self.load_model(model_name)
        self.context_manager = LongContextManager(max_context_length)
        self.tokenizer = self.load_tokenizer(model_name)
    
    def load_model(self, model_name):
        # 实际实现会加载相应的模型
        return lambda prompt: f"模型响应：{prompt}"
    
    def load_tokenizer(self, model_name):
        # 实际实现会加载相应的分词器
        return lambda text: text.split()
    
    def analyze_document(self, document_path, query):
        # 读取文档
        with open(document_path, 'r', encoding='utf-8') as f:
            document = f.read()
        
        # 添加文档到上下文
        full_context = self.context_manager.add_text(document, self.tokenizer)
        
        # 获取相关上下文
        relevant_context = self.context_manager.get_relevant_context(query, self.tokenizer)
        
        # 构建提示
        prompt = f"""请基于以下上下文回答问题：

{relevant_context}

问题：{query}

请提供详细、准确的回答。"""
        
        # 模型推理
        response = self.model(prompt)
        
        return response
    
    def batch_analyze(self, document_paths, queries):
        results = []
        for document_path, query in zip(document_paths, queries):
            result = self.analyze_document(document_path, query)
            results.append(result)
        return results

# 示例使用
if __name__ == "__main__":
    analyzer = LongContextAnalyzer()
    response = analyzer.analyze_document(
        "large_contract.txt",
        "合同中关于违约责任的条款是什么？"
    )
    print(response)

import os
import time
from typing import List, Dict, Any

class LongContextProcessor:
    def __init__(self, model_provider="gemini", max_context_length=1000000):
        self.model_provider = model_provider
        self.max_context_length = max_context_length
        self.context = []
        self.token_count = 0
        self.tokenizer = self._get_tokenizer()
    
    def _get_tokenizer(self):
        """获取分词器"""
        if self.model_provider == "gemini":
            # 模拟Gemini的分词器
            return lambda text: text.split()
        elif self.model_provider == "gpt":
            # 模拟GPT的分词器
            return lambda text: text.split()
        else:
            # 默认分词器
            return lambda text: text.split()
    
    def add_document(self, document: str) -> bool:
        """添加文档到上下文"""
        tokens = self.tokenizer(document)
        doc_token_count = len(tokens)
        
        # 检查是否超过最大长度
        if self.token_count + doc_token_count > self.max_context_length:
            # 尝试移除最早的文档
            while self.token_count + doc_token_count > self.max_context_length:
                if not self.context:
                    return False  # 无法添加
                removed_doc = self.context.pop(0)
                removed_tokens = self.tokenizer(removed_doc)
                self.token_count -= len(removed_tokens)
        
        # 添加新文档
        self.context.append(document)
        self.token_count += doc_token_count
        return True
    
    def get_context(self, query: str, max_relevant_length=200000) -> str:
        """获取与查询相关的上下文"""
        if not self.context:
            return ""
        
        # 计算每个文档与查询的相关性
        query_tokens = self.tokenizer(query)
        relevance_scores = []
        
        for i, doc in enumerate(self.context):
            doc_tokens = self.tokenizer(doc)
            score = self._calculate_relevance(query_tokens, doc_tokens)
            relevance_scores.append((i, score))
        
        # 按相关性排序
        relevance_scores.sort(key=lambda x: x[1], reverse=True)
        
        # 收集相关文档
        relevant_docs = []
        current_length = 0
        
        for i, score in relevance_scores:
            doc = self.context[i]
            doc_tokens = self.tokenizer(doc)
            doc_length = len(doc_tokens)
            
            if current_length + doc_length <= max_relevant_length:
                relevant_docs.append(doc)
                current_length += doc_length
            else:
                # 截断文档
                remaining_length = max_relevant_length - current_length
                if remaining_length > 0:
                    truncated_doc = " ".join(doc_tokens[:remaining_length])
                    relevant_docs.append(truncated_doc)
                break
        
        return " \n".join(relevant_docs)
    
    def _calculate_relevance(self, query_tokens: List[str], doc_tokens: List[str]) -> float:
        """计算查询与文档的相关性"""
        query_set = set(query_tokens)
        doc_set = set(doc_tokens)
        intersection = query_set.intersection(doc_set)
        return len(intersection) / len(query_set) if query_set else 0
    
    def process_query(self, query: str) -> str:
        """处理查询"""
        # 获取相关上下文
        relevant_context = self.get_context(query)
        
        # 构建提示
        prompt = f"""请基于以下上下文回答问题：

{relevant_context}

问题：{query}

请提供详细、准确的回答。"""
        
        # 调用模型
        response = self._call_model(prompt)
        return response
    
    def _call_model(self, prompt: str) -> str:
        """调用模型"""
        # 模拟模型调用
        time.sleep(1)  # 模拟处理时间
        return f"模型响应：{prompt[:100]}..."

# 示例使用
if __name__ == "__main__":
    processor = LongContextProcessor()
    
    # 添加文档
    processor.add_document("这是一份很长的文档，包含了很多信息..." * 1000)
    processor.add_document("这是另一份文档，包含了更多信息..." * 1000)
    
    # 处理查询
    response = processor.process_query("文档中关于X的信息是什么？")
    print(response)

class EnterpriseLongContextSystem:
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.processor = LongContextProcessor(
            model_provider=config.get("model_provider", "gemini"),
            max_context_length=config.get("max_context_length", 1000000)
        )
        self.document_store = self._init_document_store()
        self.logging = self._init_logging()
    
    def _init_document_store(self):
        """初始化文档存储"""
        # 实际实现会连接到文档数据库
        return {}
    
    def _init_logging(self):
        """初始化日志系统"""
        # 实际实现会配置日志
        return lambda msg: print(msg)
    
    def add_document(self, document_id: str, document: str) -> bool:
        """添加文档"""
        try:
            success = self.processor.add_document(document)
            if success:
                self.document_store[document_id] = document
                self.logging(f"文档 {document_id} 添加成功")
            return success
        except Exception as e:
            self.logging(f"添加文档失败: {str(e)}")
            return False
    
    def process_query(self, query: str, user_id: str) -> Dict[str, Any]:
        """处理查询"""
        start_time = time.time()
        
        try:
            response = self.processor.process_query(query)
            end_time = time.time()
            
            result = {
                "response": response,
                "processing_time": end_time - start_time,
                "user_id": user_id,
                "timestamp": time.time()
            }
            
            self.logging(f"查询处理完成，耗时: {end_time - start_time:.2f}秒")
            return result
        except Exception as e:
            self.logging(f"查询处理失败: {str(e)}")
            return {
                "error": str(e),
                "user_id": user_id,
                "timestamp": time.time()
            }
    
    def batch_process(self, queries: List[str], user_id: str) -> List[Dict[str, Any]]:
        """批量处理查询"""
        results = []
        for query in queries:
            result = self.process_query(query, user_id)
            results.append(result)
        return results

# 示例使用
if __name__ == "__main__":
    config = {
        "model_provider": "gemini",
        "max_context_length": 1000000
    }
    
    system = EnterpriseLongContextSystem(config)
    
    # 添加文档
    system.add_document("doc1", "这是一份企业合同..." * 500)
    system.add_document("doc2", "这是一份财务报告..." * 500)
    
    # 处理查询
    result = system.process_query("合同中关于付款条款的内容是什么？", "user1")
    print(result)