xy1565838851 commented 3 months ago

from monkey_model.modeling_textmonkey import TextMonkeyLMHeadModel from monkey_model.tokenization_qwen import QWenTokenizer from monkey_model.configuration_monkey import MonkeyConfig

if name == "main": checkpoint_path = "/nas_works/408972/LLM/Monkey/Monkey-Chat" input_image = "0.jpg" input_str = "Read all the text in the image." device_map = "cuda"

Create model

config = MonkeyConfig.from_pretrained(
        checkpoint_path,
        trust_remote_code=True,
    )
model = TextMonkeyLMHeadModel.from_pretrained(checkpoint_path,
    config=config,
    device_map=device_map, trust_remote_code=True).eval()
tokenizer = QWenTokenizer.from_pretrained(checkpoint_path,
                                            trust_remote_code=True)
tokenizer.padding_side = 'left'
tokenizer.pad_token_id = tokenizer.eod_id
tokenizer.IMG_TOKEN_SPAN = 1024
# tokenizer.IMG_TOKEN_SPAN = config.visual["n_queries"]

input_str = f"<img>{input_image}</img> {input_str}"
input_ids = tokenizer(input_str, return_tensors='pt', padding='longest')

attention_mask = input_ids.attention_mask
input_ids = input_ids.input_ids

pred = model.generate(
input_ids=input_ids.cuda(),
attention_mask=attention_mask.cuda(),
do_sample=True,
num_beams=1,
max_new_tokens=32768,
# max_new_tokens=2048,
min_new_tokens=1024,
length_penalty=1,
num_return_sequences=1,
output_hidden_states=True,
use_cache=True,
pad_token_id=tokenizer.eod_id,
eos_token_id=tokenizer.eod_id,
)
response = tokenizer.decode(pred[0][input_ids.size(1):].cpu(), skip_special_tokens=False).strip()
print(f"Response:{response}")