tiny-random-vopt-clip / make_tiny_model.py

data

878fa94 over 1 year ago

No virus

3.79 kB

	#!/usr/bin/env python

	# This script creates a super tiny model that is useful inside tests, when we just want to test that
	# the machinery works, without needing to check the quality of the outcomes.
	#
	# usage: adjust the configs if wanted, but otherwise just run the script

	from pathlib import Path
	from types import SimpleNamespace

	import torchvision.transforms as transforms
	from PIL import Image

	from m4.models.vopt.modeling_vopt import VOPTConfig, VOPTForCausalLM
	from m4.training.packing import image_attention_mask_for_packed_input_ids, incremental_to_binary_attention_mask
	from m4.training.utils import get_tokenizer


	mname_tiny = "tiny-random-vopt-clip"

	path = Path(mname_tiny)
	path.mkdir(parents=True, exist_ok=True)

	# from the hardcoded https://github.com/huggingface/m4/blob/adf102f0000cb2632cd8a3ebb87398c65e448a97/m4/training/main.py#L80
	additional_vocab_size = 2

	config = VOPTConfig()
	config.update(
	dict(
	ffn_dim=64,
	hidden_size=16,
	max_position_embeddings=128,
	num_attention_heads=4,
	num_hidden_layers=2,
	word_embed_proj_dim=16,
	max_new_tokens=100,
	use_resampler=True,
	resampler_depth=2,
	resampler_head_dim=8,
	resampler_n_heads=2,
	resampler_n_latents=16,
	vision_embed_dim=32,
	vision_image_size=30,
	vision_model_name="hf-internal-testing/tiny-random-clip",
	vision_model_params="{}",
	vocab_size=50265,
	additional_vocab_size=additional_vocab_size,
	)
	)

	# print(config)
	# can now modify config to say tiny values

	model = VOPTForCausalLM.from_config(config)
	# print(model.config)
	# print(model)

	tokenizer_config = dict(
	tokenizer_add_special_tokens="{}",
	tokenizer_add_tokens=(
	'[AddedToken("<fake_token_around_image>", rstrip=False, lstrip=False), AddedToken("<image>", rstrip=False,'
	" lstrip=False)]"
	),
	tokenizer_name="facebook/opt-13b",
	tokenizer_params='{"use_fast":True}',
	)
	tokenizer_config = SimpleNamespace(**tokenizer_config)
	# print(tokenizer_config)

	tokenizer = get_tokenizer(
	tokenizer_name=tokenizer_config.tokenizer_name,
	tokenizer_add_tokens=tokenizer_config.tokenizer_add_tokens,
	tokenizer_add_special_tokens=tokenizer_config.tokenizer_add_special_tokens,
	tokenizer_params=tokenizer_config.tokenizer_params,
	additional_vocab_size=model.config.additional_vocab_size,
	model_vocab_size=model.config.vocab_size,
	)
	assert "<image>" in tokenizer.get_vocab()

	# Test w/ one image and one text
	query = "<fake_token_around_image><image><fake_token_around_image>This is a picture of a cat."
	query_tokens = tokenizer(query, return_tensors="pt")

	num_images_per_ex = 1
	pixel_values = transforms.ToTensor()(Image.new("RGB", (30, 30))).repeat(1, 1, 1, 1).unsqueeze(0)
	image_attention_mask, _ = image_attention_mask_for_packed_input_ids(query_tokens["input_ids"], tokenizer)
	image_attention_mask = incremental_to_binary_attention_mask(image_attention_mask, num_classes=num_images_per_ex)

	input = {
	"input_ids": query_tokens["input_ids"],
	"attention_mask": query_tokens["attention_mask"],
	"pixel_values": pixel_values,
	"pixel_values": pixel_values,
	"image_attention_mask": image_attention_mask,
	}
	# debug shapes
	# print(query_tokens["input_ids"].shape)
	# print(query_tokens["attention_mask"].shape)
	# print(pixel_values.shape)
	# print(image_attention_mask.shape)

	out_gen = model.generate(**input)
	text = tokenizer.batch_decode(out_gen)
	# print(text)

	# Save model + config + tokenizer
	model.half() # makes it smaller
	model.save_pretrained(path)
	tokenizer.save_pretrained(path)

	# test we can load it back
	model = VOPTForCausalLM.from_pretrained(path)

	print(f"Generated {mname_tiny} - Upload the generated folder to the hub")