import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig

model_name = "Qwen/Qwen2.5-7B-Instruct"

# 方法一：INT8 量化（推荐用于生产环境）
# 显存减少约 50%，推理速度基本持平，精度损失 < 1%
model_int8 = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_8bit=True,   # 启用 INT8 量化
    device_map="auto"    # 自动分配 GPU/CPU
)

# 方法二：4-bit NF4 量化（推荐用于显存受限场景）
# 显存减少约 75%，精度损失 1-3%，是 QLoRA 微调时使用的格式
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",               # NF4 格式，精度优于普通 int4
    bnb_4bit_compute_dtype=torch.float16     # 计算时恢复为 float16，提升精度
)
model_4bit = AutoModelForCausalLM.from_pretrained(
    model_name,
    quantization_config=bnb_config,
    device_map="auto"
)

# 方法三：GGUF 量化（用于 llama.cpp / Ollama，适合本地部署）
# 优势：支持纯 CPU 推理，适合没有 GPU 的场景
# 使用 llama.cpp 工具转换：
# python convert_hf_to_gguf.py model_dir --outtype q4_k_m --outfile model_q4.gguf

from transformers import AutoModelForCausalLM
import torch

# 安装：pip install flash-attn --no-build-isolation
# 要求：NVIDIA GPU Ampere 架构及以上（A100/A10/RTX 30xx+）

# 适合场景：处理长文档、长对话历史（2K+ token）
model = AutoModelForCausalLM.from_pretrained(
    "Qwen/Qwen2.5-7B-Instruct",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",  # 启用 Flash Attention 2
    device_map="auto"
)

# 验证是否启用成功
print(model.config._attn_implementation)  # 应输出 "flash_attention_2"

import torch
import torch.nn as nn
import torch.nn.utils.prune as prune

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(128, 64)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(64, 10)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

model = SimpleNet()

# 结构化剪枝：按 L1 范数删除整行（对应整个输出神经元）
# amount=0.3 表示删除 30% 的输出神经元
prune.ln_structured(
    model.fc1,
    name="weight",
    amount=0.3,
    n=1,        # L1 范数
    dim=0       # 按行（输出维度）剪枝
)

# 查看剪枝效果
weight = model.fc1.weight
sparsity = (weight == 0).float().mean().item()
print(f"fc1 层权重稀疏度: {sparsity:.1%}")

# 永久化剪枝（实际删除参数）
prune.remove(model.fc1, "weight")

total_params = sum(p.numel() for p in model.parameters())
nonzero_params = sum(p.nonzero().shape[0] for p in model.parameters())
print(f"总参数量: {total_params}")
print(f"非零参数量: {nonzero_params}")
print(f"整体稀疏度: {1 - nonzero_params/total_params:.1%}")

import torch
import torch.onnx
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model_name = "bert-base-chinese"
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model.eval()

# 准备示例输入
sample_text = "今天天气很好"
inputs = tokenizer(sample_text, return_tensors="pt", max_length=128, padding="max_length")

# 导出为 ONNX 格式
torch.onnx.export(
    model,
    (inputs['input_ids'], inputs['attention_mask']),
    "bert_classifier.onnx",
    input_names=['input_ids', 'attention_mask'],
    output_names=['logits'],
    dynamic_axes={                           # 允许动态 batch size 和序列长度
        'input_ids': {0: 'batch_size', 1: 'sequence'},
        'attention_mask': {0: 'batch_size', 1: 'sequence'},
        'logits': {0: 'batch_size'}
    },
    opset_version=14
)
print("ONNX 模型已导出到 bert_classifier.onnx")

# 用 ONNX Runtime 进行推理（pip install onnxruntime）
import onnxruntime as ort
import numpy as np

session = ort.InferenceSession(
    "bert_classifier.onnx",
    providers=['CPUExecutionProvider']  # CPU 推理
)

inputs_onnx = {
    'input_ids': inputs['input_ids'].numpy(),
    'attention_mask': inputs['attention_mask'].numpy()
}
outputs = session.run(None, inputs_onnx)
print(f"ONNX 推理结果维度: {outputs[0].shape}")