File size: 5,549 Bytes
a512fa8 36ecd25 a512fa8 | 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 99 100 101 102 103 104 105 106 107 | from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from zipfile import ZipFile
import torch
from safetensors.torch import load_file
from transformers import AutoConfig,AutoProcessor
from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Model,Qwen3_5VisionModel
from inference import Engine,Stream
class VisualModel(torch.nn.Module):
get_rope_index=Qwen3_5Model.get_rope_index
get_vision_position_ids=Qwen3_5Model.get_vision_position_ids
get_placeholder_mask=Qwen3_5Model.get_placeholder_mask
def __init__(self,text,config,visual):
super().__init__()
self.text,self.config,self.visual=text,config,visual
def get_input_embeddings(self):
return self.text.get_input_embeddings()
def forward(self,input_ids,attention_mask,pixel_values,image_grid_thw):
embeds=self.get_input_embeddings()(input_ids)
features=self.visual(pixel_values.to(dtype=embeds.dtype),grid_thw=image_grid_thw,return_dict=True).pooler_output
masks=self.get_placeholder_mask(input_ids,embeds,image_features=features)
embeds=embeds.masked_scatter(masks[0],features.to(embeds))
rope,deltas=self.get_rope_index(input_ids,mm_token_type_ids=(input_ids==self.config.image_token_id).long(),
image_grid_thw=image_grid_thw,attention_mask=attention_mask)
physical=(attention_mask.long().cumsum(-1)-1).clamp_min(0)
positions=torch.cat((physical.unsqueeze(0),rope),dim=0)
output=self.text(inputs_embeds=embeds,attention_mask=attention_mask,position_ids=positions,use_cache=True,logits_to_keep=1)
lengths=attention_mask.sum(-1)
spatial=lengths+deltas.flatten()
output['vision_decode_positions']=torch.stack((lengths,spatial,spatial,spatial))[:,:,None]
return output
class VisualContext:
def __init__(self,model,stream,batch):
self.model,self.stream=model,stream
self.positions=torch.empty((4,2*batch+1,1),dtype=torch.long,device='cuda')
def __call__(self,ids,mask,indices):
selected=[self.stream.rows[i]['vision_inputs'] for i in indices.tolist()]
output=self.model(input_ids=ids,attention_mask=mask,
pixel_values=torch.cat([v['pixel_values'] for v in selected]).to(ids.device),
image_grid_thw=torch.cat([v['image_grid_thw'] for v in selected]).to(ids.device))
self.positions.index_copy_(1,indices,output['vision_decode_positions'])
return output
def decode_positions(self,indices):
return self.positions.index_select(1,indices)
def generate_visual(args,claim,write,report):
engine=Engine(args.model,args.batch,args.profile)
config=AutoConfig.from_pretrained(args.model/'vision',local_files_only=True)
config.vision_config._attn_implementation='flash_attention_2'
visual=Qwen3_5VisionModel(config.vision_config)
visual.load_state_dict(load_file(str(args.model/'vision/model.safetensors')))
model=VisualModel(engine.model,config,visual).cuda().to(dtype=torch.bfloat16).eval().requires_grad_(False)
processor=AutoProcessor.from_pretrained(args.model/'vision',local_files_only=True)
processor.tokenizer.padding_side='left'
processor.image_processor.size={'longest_edge':4194304,'shortest_edge':65536}
media=ZipFile(args.model/'evaluation_data/vision/media.zip')
pool=ThreadPoolExecutor(20)
def image(row):
from PIL import Image
with media.open(row['image_path']) as source, Image.open(source) as image:
return processor.image_processor(images=image.convert('RGB'),return_tensors='pt')
def fetch(count):
rows=claim(count)
if not rows:return []
images=list(pool.map(image,rows))
texts=[]
for row,pixels in zip(rows,images):
question=((row.get('hint') or '')+'\n'+row['question']).strip()
choices={k:v for k,v in (row.get('choices') or {}).items() if v is not None}
row['choices']=choices
if choices:
question+='\n'+'\n'.join(f'{k}. {v}' for k,v in choices.items())
question+='\nGive only the option letter as your final answer.'
elif row['source'] in ('mme','pope'):
question+='\nGive only Yes or No as your final answer.'
else:question+='\nGive a short answer without explanation as your final answer.'
text=processor.apply_chat_template([{'role':'user','content':[{'type':'image'},{'type':'text','text':question}]}],
tokenize=False,add_generation_prompt=True,enable_thinking=True)
count=int(pixels['image_grid_thw'].prod())//processor.image_processor.merge_size**2
texts.append(text.replace(processor.image_token,processor.image_token*count,1))
ids=processor.tokenizer(texts,add_special_tokens=False,padding=False)['input_ids']
return [dict(row,input_ids=token,vision_inputs=pixels,max_tokens=args.max_new_tokens)
for row,token,pixels in zip(rows,ids,images)]
first=fetch(args.batch)
if not first:return {}
stream=Stream(engine,fetch,write,report)
stream.prefill_chunk_size=lambda width:4
context=VisualContext(model,stream,args.batch)
try:
return engine.run_stream(stream,first,cap=args.max_new_tokens,temperature=1,top_k=20,top_p=.95,
presence_penalty=1.5,seed=args.seed,mechanical=False,prefill_context=context,
eos_ids=(processor.tokenizer.eos_token_id,))
finally:
pool.shutdown(wait=True)
media.close()
|