File size: 2,439 Bytes
5076284 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 |
import json
from transformers import AutoTokenizer
import torch
import llava.model.language_model.llava_olmo1p58b as llava_olmo ##
import llava.model.language_model.llava_llama as llava_llama
from OLMo_Bitnet_1B.modeling_olmo import OLMoForCausalLM
from PIL import Image
import requests
from llava.mm_utils import tokenizer_image_token, process_images, get_model_name_from_path
from llava.conversation import conv_templates
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
DEFAULT_IMAGE_TOKEN = "<image>"
IMAGE_TOKEN_INDEX = -200
# Define Image and Text inputs..
text = "What are the four major tournaments of the sport shown in the image?"
url = "https://farm3.staticflickr.com/2157/2439959136_d932f4e816_z.jpg"
image = Image.open(requests.get(url, stream=True).raw)
# LOAD MODEL FROM CHECKPOINT
with open('./checkpoints/llava-LlavaOLMoBitnet1B-Run3-finetune/config.json') as json_file:
data = json.load(json_file)
config_class = llava_olmo.LlavaOLMoBitnet1BConfig(**data)
model = llava_olmo.LlavaOLMoBitnet1BForCausalLM(config_class).to(device)
weight_checkpoint = torch.load('./checkpoints/llava-LlavaOLMoBitnet1B-Run3-finetune/pytorch_model.bin')
model.load_state_dict(weight_checkpoint)
# pre-process image; Apply chat template and tokenize text
image_processor = model.model.vision_tower.image_processor
tokenizer = AutoTokenizer.from_pretrained(
"NousResearch/OLMo-Bitnet-1B",
model_max_length=2048,
padding_side="right",
pad_token_id=1,
use_fast=True,
legacy=False,
unk_token='<|padding|>',
)
image_tensor = process_images([image], image_processor, model.config)[0]
text = DEFAULT_IMAGE_TOKEN + '\n' + text
conv = conv_templates['llava_v1'].copy()
conv.append_message(conv.roles[0], text)
conv.append_message(conv.roles[1], None)
prompt = conv.get_prompt()
text_tokens = tokenizer_image_token(prompt, tokenizer, IMAGE_TOKEN_INDEX, return_tensors='pt').unsqueeze(0).to(device)
# Generate response from the model
response = model.generate(images=image_tensor.unsqueeze(0).to(device), inputs=text_tokens, max_new_tokens=400)
decoded_text = tokenizer.batch_decode(response, skip_special_tokens=True)[0]
print("\n\n", "-"*100)
print(decoded_text[:decoded_text.find('</s>')].replace('|||IP_ADDRESS|||', '')) # The replace part is due to unwanted token introduction at start
print("-"*100)
|