#!/usr/bin/env python3
import torch
from typing import Optional

def compress_context(context: torch.Tensor, mode: str = "balanced"):
    """Compress context tensor."""
    original_len = context.shape[0]
    target_ratio = 0.5 if mode == "balanced" else 0.6
    target_len = int(original_len * target_ratio)
    
    norms = torch.norm(context, dim=-1)
    _, top_indices = torch.topk(norms, target_len)
    top_indices = torch.sort(top_indices).values
    
    compressed = context[top_indices]
    return compressed

if __name__ == "__main__":
    device = "cuda" if torch.cuda.is_available() else "cpu"
    context = torch.randn(2048, 768, device=device)
    result = compress_context(context)
    print(f"Compressed: {context.shape[0]} -> {result.shape[0]}")
