File size: 3,850 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 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
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)
#
##
#
#
#
'''
# ORIGINAL CODE WITH ONLY OLMO:
with open('llava/config.json') as json_file:
data = json.load(json_file)
text = "Paris is a historic city with architectural marvels. It is also "
# text = ["Language modeling is "]
config_class = llava_olmo.LlavaOLMoBitnet1BConfig(**data)
lolmo = llava_olmo.LlavaOLMoBitnet1BForCausalLM(config_class).to(device)
lolmo.load_state_dict(torch.load('OLMo_Bitnet_1B/pytorch_model.bin'), strict=False)
olmo = OLMoForCausalLM(config_class).to(device)
olmo.load_state_dict(torch.load('OLMo_Bitnet_1B/pytorch_model.bin'))
actual_olmo = OLMoForCausalLM.from_pretrained("allenai/OLMo-1B").to(device)
actual_olmo_tokenizer = OLMoTokenizerFast.from_pretrained("allenai/OLMo-1B")
olmo_tokenizer = AutoTokenizer.from_pretrained("NousResearch/OLMo-Bitnet-1B")
olmo_tokens = olmo_tokenizer(text, return_tensors='pt', return_token_type_ids=False).to(device)
# olmo_tokens = actual_olmo_tokenizer(text, return_tensors='pt', return_token_type_ids=False).to(device)
response = lolmo.generate(inputs=olmo_tokens['input_ids'], attention_mask=olmo_tokens['attention_mask'], max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
# response = olmo.generate(inputs=olmo_tokens['input_ids'], attention_mask=olmo_tokens['attention_mask'], max_new_tokens=100, do_sample=True, top_k=50, top_p=0.95)
print(olmo_tokenizer.batch_decode(response, skip_special_tokens=True)[0])
''' |