import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from shortx import AutoShortGPT, PruningStrategy

# Initialize model with automatic adapter selection
model = AutoShortGPT.from_pretrained("Qwen/Qwen2.5-7B-Instruct")
# Automatically selects QwenShortGPT adapter for Qwen models
# or LlamaShortGPT adapter for Llama models

# Analyze layer importance
analysis = model.analyze_layers()
print(f"Layer importance scores: {analysis['importances']}")
print(f"Recommended layers to prune: {analysis['recommended_prune_layers']}")

# Apply pruning with non-invasive patching
with model.enable_monkey_patching() as m:
    # Skip layers completely using IDENTITY strategy
    m.prune_layers_patch([27, 26, 25, 28, 24, 29, 23, 21, 22], PruningStrategy.IDENTITY)
    model_inputs = m.tokenizer(["Hello world"], return_tensors="pt").to(m.model.device)
    # Generate with pruned model
    result_ids = m.model.generate(**model_inputs, max_length=100)
    result = m.tokenizer.decode(result_ids[0, model_inputs.input_ids.shape[-1]:])
    print(result)
# Model automatically restored to original state

# Mixed pruning strategies for fine-grained control
layer_config = {
    10: PruningStrategy.IDENTITY,       # Skip completely
    15: PruningStrategy.ATTENTION_ONLY, # Keep only attention
    20: PruningStrategy.FFN_ONLY        # Keep only FFN
}
model.prune_layers_mixed(layer_config)
result = model.generate("The meaning of life is", max_length=50)

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
from shortx import AutoShortV

# Initialize ShortV with automatic adapter selection
shortv = AutoShortV.from_pretrained(
    model_name="llava-hf/llava-1.5-7b-hf",
    torch_dtype="auto",
    device_map="cuda:0"
)

# Enable layer skipping optimization
shortv.enable_monkey_patching()

# Prepare conversation with image
conversation = [
    {
        "role": "user",
        "content": [
            {"type": "image", "url": "examples/cartoon_image.jpg"},
            {"type": "text", "text": "What is shown in this image?"},
        ],
    },
]

# Process inputs
inputs = shortv.processor.apply_chat_template(
    conversation,
    add_generation_prompt=True,
    tokenize=True,
    return_dict=True,
    return_tensors="pt"
).to(shortv.model.device)

# Generate
generate_ids = shortv.model.generate(**inputs, max_new_tokens=100)
decoded = shortv.processor.decode(generate_ids[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True)
print(f"Response: {decoded.strip()}")

# Analyze layer contribution for optimization
results = shortv.calculate_lc_scores()
print(f"Layer importance ranking: {results['layer_ranking_by_importance']}")
print(f"Recommended layers to skip: {results['recommended_skip_layers']}")

# Dynamically change skip layers
shortv.set_skip_layers([0, 2, 4, 6, 8, 10, 12, 14, 16, 18, 20])

# Generate
generate_ids = shortv.model.generate(**inputs, max_new_tokens=100)
decoded = shortv.processor.decode(generate_ids[0, inputs.input_ids.shape[-1]:], skip_special_tokens=True)
print(f"Response: {decoded.strip()}")

# Disable layer skipping to restore original model
shortv.disable_monkey_patching()