from transformers import AutoProcessor, PaliGemmaConfig, PaliGemmaForConditionalGeneration
config = PaliGemmaConfig.from_pretrained("google/paligemma-3b-pt-224")
config.text_config.intermediate_size = 16
config.text_config.num_hidden_layers = 2
config.text_config.hidden_size = 16
config.vision_config.num_hidden_layers = 16
config.vision_config.num_hidden_layers = 2
config.vision_config.hidden_size = 16
config.vision_config.projection_dim = 16
config.vision_config.intermediate_size = 8
model = PaliGemmaForConditionalGeneration(config=config)
processor = AutoProcessor.from_pretrained("google/paligemma-3b-pt-224")
processor.chat_template = """{% if not add_generation_prompt is defined %}{% set add_generation_prompt = false %}{% endif %}{% for message in messages %}<|im_start|>{% if message['role'] == 'user' %}USER: {% else %}ASSISTANT: {% endif %}{% for item in message['content'] if item['type'] == 'text' %}{{ item['text'] }}<|im_end|>{% endfor %}{% if message['role'] == 'user' %} {% else %}{{eos_token}}{% endif %}{% endfor %}{% if add_generation_prompt %}ASSISTANT: {% endif %}"""
model_id = "trl-internal-testing/tiny-random-paligemma"
model.push_to_hub(model_id)
processor.push_to_hub(model_id)
