# 推理的重要性
inference_importance = {
    "应用核心": "推理是模型应用的核心环节",
    "性能影响": "推理性能直接影响用户体验",
    "成本控制": "优化推理可以降低成本",
    "实时性": "推理速度影响实时应用"
}

# 文本生成推理流程
def text_generation_inference(model, prompt, max_length=100):
    """
    文本生成推理
    """
    # 1. 输入编码
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    
    # 2. 模型推理
    outputs = model.generate(
        input_ids,
        max_length=max_length,
        num_return_sequences=1,
        temperature=0.7,
        do_sample=True
    )
    
    # 3. 输出解码
    generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    
    # 4. 返回结果
    return generated_text

# Temperature采样
def temperature_sampling(logits, temperature=1.0):
    """
    temperature: 温度参数
    - temperature < 1.0: 更确定，更保守
    - temperature = 1.0: 标准采样
    - temperature > 1.0: 更随机，更创新
    """
    # 应用temperature
    scaled_logits = logits / temperature
    
    # Softmax
    probs = softmax(scaled_logits)
    
    # 采样
    token = sample(probs)
    
    return token

# Top-k采样
def top_k_sampling(logits, k=50):
    """
    k: 只考虑概率最高的k个词
    """
    # 获取top-k
    top_k_logits, top_k_indices = torch.topk(logits, k)
    
    # 计算概率
    probs = softmax(top_k_logits)
    
    # 采样
    sampled_index = sample(probs)
    token = top_k_indices[sampled_index]
    
    return token

# Top-p采样 (核采样)
def top_p_sampling(logits, p=0.9):
    """
    p: 累积概率阈值（通常0.9-0.95）
    """
    # 排序
    sorted_logits, sorted_indices = torch.sort(logits, descending=True)
    sorted_probs = softmax(sorted_logits)
    
    # 计算累积概率
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
    
    # 找到累积概率达到p的位置
    sorted_indices_to_remove = cumulative_probs > p
    sorted_indices_to_remove[..., 0] = False  # 至少保留一个
    
    # 过滤
    indices_to_remove = sorted_indices_to_remove.scatter(
        1, sorted_indices, sorted_indices_to_remove
    )
    logits[indices_to_remove] = float('-inf')
    
    # 采样
    probs = softmax(logits)
    token = sample(probs)
    
    return token

# 基础推理示例
from transformers import AutoModelForCausalLM, AutoTokenizer

# 加载模型和分词器
model_name = "gpt2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)

# 推理
prompt = "The future of AI is"
input_ids = tokenizer.encode(prompt, return_tensors="pt")

# 生成
outputs = model.generate(
    input_ids,
    max_length=100,
    num_return_sequences=1,
    temperature=0.7,
    top_p=0.9,
    do_sample=True
)

# 解码
generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

# 流式推理示例
def stream_generation(model, tokenizer, prompt):
    """
    流式生成，实时输出token
    """
    input_ids = tokenizer.encode(prompt, return_tensors="pt")
    
    # 流式生成
    for output in model.generate(
        input_ids,
        max_length=100,
        do_sample=True,
        streamer=None  # 可以设置streamer
    ):
        # 解码并输出
        token = tokenizer.decode([output], skip_special_tokens=True)
        print(token, end="", flush=True)

# 批量推理示例
def batch_inference(model, tokenizer, prompts, batch_size=8):
    """
    批量推理
    """
    results = []
    
    # 分批处理
    for i in range(0, len(prompts), batch_size):
        batch_prompts = prompts[i:i+batch_size]
        
        # 编码
        inputs = tokenizer(
            batch_prompts,
            return_tensors="pt",
            padding=True,
            truncation=True
        )
        
        # 生成
        outputs = model.generate(
            **inputs,
            max_length=100,
            do_sample=True
        )
        
        # 解码
        generated_texts = tokenizer.batch_decode(
            outputs,
            skip_special_tokens=True
        )
        
        results.extend(generated_texts)
    
    return results