Huihui-gemma-3n-E2B-it-abliterated

from transformers import AutoProcessor, Gemma3nForConditionalGeneration
from transformers.models.gemma3n.modeling_gemma3n import Gemma3nTextDecoderLayer
from PIL import Image
import requests
import torch
import jaxtyping
import einops
import base64

model_id = "google/gemma-3n-E2B-it"

model = Gemma3nForConditionalGeneration.from_pretrained(model_id, device_map="cuda", torch_dtype=torch.bfloat16,).eval()

#refusal_dir= torch.load(model_id + "/final_refusal_dirs-16.pt", map_location='cpu', weights_only=True)

#refusal_dir = refusal_dir.to(model.device)
#refusal_dir = refusal_dir.to(torch.bfloat16)
#
#def direction_ablation_hook(activation: jaxtyping.Float[torch.Tensor, "... d_act"],
#                            direction: jaxtyping.Float[torch.Tensor, "d_act"]):
#    proj = einops.einsum(activation, direction.view(-1, 1), '... d_act, d_act single -> ... single') * direction
#    return activation - proj
#
#class AblationDecoderLayer(Gemma3nTextDecoderLayer):
#    def __init__(self, original_layer, config, layer_idx, refusal_dir):
#        super(AblationDecoderLayer, self).__init__(config, layer_idx)
#        self.original_layer = original_layer
#        self.refusal_dir = refusal_dir
#
#    def forward(self, *args, **kwargs):
#        hidden_states = args[0]
#        ablated = direction_ablation_hook(hidden_states, self.refusal_dir.to(hidden_states.device)).to(hidden_states.device)
#        args = (ablated,) + args[1:]
#        return self.original_layer.forward(*args, **kwargs)
#
#for idx in range(len(model.model.language_model.layers)):
#    model.model.language_model.layers[idx] = AblationDecoderLayer(model.model.language_model.layers[idx], model.config.text_config, idx, refusal_dir)

processor = AutoProcessor.from_pretrained(model_id)

# https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/bee.jpg
image_path = model_id + "/bee.jpg"

with open(image_path, "rb") as image_file:
    encoded_string = base64.b64encode(image_file.read()).decode("utf-8")
    image_type = "image/jpeg" if image_path.endswith(".jpg") else "image/png"
    data_uri = f"data:{image_type};base64,{encoded_string}"

    messages = [
        {
            "role": "system",
            "content": [{"type": "text", "text": "You are a helpful assistant."}]
        },
        {
            "role": "user",
            "content": [
                {"type": "image", "image": data_uri},
                {"type": "text", "text": "Describe this image in detail."}
            ]
        }
    ]

    inputs = processor.apply_chat_template(
        messages,
        add_generation_prompt=True,
        tokenize=True,
        return_dict=True,
        return_tensors="pt",
    ).to(model.device, dtype=torch.bfloat16)

    input_len = inputs["input_ids"].shape[-1]

    with torch.inference_mode():
        generation = model.generate(**inputs, max_new_tokens=1024, do_sample=True)
        generation = generation[0][input_len:]

    decoded = processor.decode(generation, skip_special_tokens=True)
    print(decoded)


    # ## A Vibrant Garden Scene: A Close-Up of a Busy Bee
    # 
    # This image is a delightful close-up shot showcasing a charming bee in full swing, visiting a beautiful daisy-like flower in a garden. Here's a detailed breakdown:
    # 
    # **Overall Composition:**
    # 
    # The image is divided into three main sections: a large central focus on a pink flower, a smaller top section indicating another pink flower, and a bottom right section highlighting a vibrant red flower. This creates a sense of a bustling garden scene.
    # 
    # **Central Focus: The Mighty Bee on a Pink Daisy:**
    # 
    # * **The Bee:** The star of the show is a plump, fuzzy bee positioned prominently on a bright pink daisy.
    #     * **Body:** The bee has a dark, almost black body with a yellow band across its abdomen. Its wings are visible, slightly spread, giving it a charming "flying" look.
    #     * **Wings:**  It has a distinctive crown of small, white and black stripes.
    #     * **Position:** The bee is centered on a large, open flower head, often called a daisy, making it easily identifiable.
    # * **The Daisy:**
    #     * **Color:** A vibrant pink color dominates the central area.
    #     * **Petals:** The daisy has four distinct, rounded petals, each with a subtle inner line for definition.
    #     * **Size:** It's a generous size, with a yellow center and a green background.
    #     * **Variety:** A smaller version of the same daisy is placed above, allowing for easy comparison.
    # 
    # **Background Details:**
    # 
    # * **Background:** The background is divided into two rows of flowers, creating a visually appealing layering effect.
    #     * **Upper Row (Larger):**  A large, bright pink flower dominates the top half of the image.
    #     * **Lower Row (Smaller):**  A series of smaller flowers are arranged for a more detailed view.
    # * **Flower Types:**
    #     * **Red Flower:** A bold red flower is positioned towards the bottom right, with a small white dot on the petals for added visual interest.
    #     * **Smaller Pink Flowers:**  Several smaller pink flowers are arranged around the main crop, creating a sense of a dense garden.
    #     * **Brown Buds:** A cluster of brown buds is placed towards the bottom left, with a smaller pink flower peeking out.
    # 
    # **Overall Style:**
    # 
    # The image has a charming, slightly rustic feel, often used for illustrating items or creating a consistent visual style. The use of a light background makes the flowers and bee easily legible.
    # 
    # **In summary:**
    # 
    # This is a delightful image highlighting the busy life of a bee in a thriving garden. It's a vibrant scene, well-organized, and easily readable, making it perfect for any "garden" enthusiast!