flax-community
/

vit-gpt2

Model card Files Files and versions Metrics Training metrics Community

vit-gpt2 / test_model.py

ydshieh

add model.py and app.py

f338d56 almost 3 years ago

raw history blame

No virus

2.16 kB

	import sys, os

	current_path = os.path.dirname(os.path.abspath(__file__))
	sys.path.append(current_path)

	# Main model - ViTGPT2LM
	from vit_gpt2.modeling_flax_vit_gpt2_lm import FlaxViTGPT2LMForConditionalGeneration

	# Vit - as encoder
	from transformers import ViTFeatureExtractor
	from PIL import Image
	import requests
	import numpy as np

	# GPT2 / GPT2LM - as decoder
	from transformers import ViTFeatureExtractor, GPT2Tokenizer

	model_name_or_path = './outputs/ckpt_2/'
	flax_vit_gpt2_lm = FlaxViTGPT2LMForConditionalGeneration.from_pretrained(model_name_or_path)

	vit_model_name = 'google/vit-base-patch16-224-in21k'
	feature_extractor = ViTFeatureExtractor.from_pretrained(vit_model_name)

	gpt2_model_name = 'asi/gpt-fr-cased-small'
	tokenizer = GPT2Tokenizer.from_pretrained(gpt2_model_name)

	max_length = 32
	num_beams = 16
	gen_kwargs = {"max_length": max_length, "num_beams": num_beams}


	# encoder data
	url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
	image = Image.open(requests.get(url, stream=True).raw)
	# batch dim is added automatically
	encoder_inputs = feature_extractor(images=image, return_tensors="jax")
	pixel_values = encoder_inputs.pixel_values
	print(f'pixel_values.shape = {pixel_values.shape}')

	# decoder data
	sentence = 'mon chien est mignon'
	# IMPORTANT: For training/evaluation/attention_mask/loss
	sentence += ' ' + tokenizer.eos_token
	# batch dim is added automatically
	decoder_inputs = tokenizer(sentence, return_tensors="jax")
	print(decoder_inputs)
	print(f'input_ids.shape = {decoder_inputs.input_ids.shape}')

	# model data
	inputs = dict(decoder_inputs)
	inputs['pixel_values'] = pixel_values


	logits = flax_vit_gpt2_lm(**inputs)[0]
	preds = np.argmax(logits, axis=-1)
	print('=' * 60)
	print('Flax: Vit-GPT2-LM')
	print('predicted token ids:')
	print(preds)
	print('=' * 60)


	# Generation!
	batch = {'pixel_values': pixel_values}
	generation = flax_vit_gpt2_lm.generate(batch['pixel_values'], **gen_kwargs)
	print('generation:')
	print(generation)
	print('=' * 60)

	token_ids = np.array(generation.sequences)[0]
	caption = tokenizer.decode(token_ids)
	print(f'token_ids: {token_ids}')
	print(f'caption: {caption}')
	print('=' * 60)