import torch, os from peft import LoraConfig, get_peft_model, PeftModel from transformers import AutoModelForCausalLM, Cache from transformers.utils import logging from .tokenization_live import build_live_tokenizer_and_update_config from .vision_live import build_live_vision logger = logging.get_logger(__name__) class LiveMixin(AutoModelForCausalLM): def set_vision_inside(self): logger.warning_once("!!! Set vision encoder in the model, only recommended for on in-the-wild inference. " "Please dont call this for efficient training & evaluation. Instead, do visual feature pre-extraction.") self.vision_encoder, self.vision_encode = build_live_vision(self.config) def unset_vision_inside(self): del self.vision_encoder del self.vision_encode def visual_embed(self, frames: torch.Tensor): if hasattr(self, 'vision_encode'): with torch.cuda.amp.autocast(): frames = self.vision_encode(self.vision_encoder, frames) frames = frames.to(self.dtype) frames = self.connector(frames) return frames.view(-1, frames.shape[-1]) def joint_embed( self, input_ids: torch.Tensor = None, frames: torch.Tensor = None, ): if frames is None: return self.get_input_embeddings()(input_ids) if input_ids is None: return self.visual_embed(frames) inputs_embeds = self.get_input_embeddings()(input_ids.clamp(max=self.vocab_size-1)) v_mask = input_ids == self.config.v_placeholder_id if v_mask.any(): inputs_embeds[v_mask] = self.visual_embed(frames) return inputs_embeds @torch.no_grad() def stream_evaluate( self, input_ids: torch.LongTensor, labels: torch.LongTensor, frames: torch.ByteTensor, ignore_token_id: int = -100, frame_token_interval_threshold: float = 0.0, **kwargs ): # 0. evaluation only supports batch_size = 1 assert input_ids.size(0) == labels.size(0) == 1 input_id, label = input_ids[0], labels[0] device = input_id.device zero = torch.tensor(0, dtype=torch.int, device=device) one = torch.tensor(1, dtype=torch.int, device=device) # 1. prepare multi-turn start and stop turn_stops = ((input_id == self.config.eos_token_id).nonzero() + 1)[:,0].tolist() turn_starts = [0] + turn_stops[:-1] num_turns = len(turn_starts) # 2. forward the full input_ids and labels, get tokenwise logits and losses outputs = self.forward(input_ids=input_ids, frames=frames, return_dict=True, use_cache=True) logit, past_key_values = outputs.logits[0], outputs.past_key_values # 3. compute metrics for each turn v_placeholder_id = self.config.v_placeholder_id use_interval = self.config.frame_token_interval_id is not None frame_token_interval_id = self.config.frame_token_interval_id if use_interval else self.config.eos_token_id frame_num_tokens = self.config.frame_token_cls if self.config.frame_token_pooled: frame_num_tokens += self.config.frame_token_pooled[0] * self.config.frame_token_pooled[1] past_num_frames = 0 lm_ppls, frame_diffs, fluencies, lm_correctness = [], [], [], [] for r, (turn_start, turn_stop) in enumerate(zip(turn_starts, turn_stops)): ## 3.1. we only have two losses: stream loss on frame tokens, and lm loss. prepare corresponding mask according two losses turn_label = label[turn_start:turn_stop] turn_learn_mask = turn_label != ignore_token_id if not turn_learn_mask.any(): continue turn_logit = logit[turn_start:turn_stop] turn_input_id = input_id[turn_start:turn_stop] turn_v_mask = turn_input_id == v_placeholder_id turn_num_frames = turn_v_mask.sum() // frame_num_tokens turn_stream_mask = turn_v_mask & turn_learn_mask turn_lm_mask = turn_learn_mask & ~turn_stream_mask ## 3.2 ppl, offline metric if turn_lm_mask.any(): turn_lm_masked_logit, turn_lm_masked_label = turn_logit[turn_lm_mask], turn_label[turn_lm_mask] lm_ppl = torch.nn.functional.cross_entropy(turn_lm_masked_logit, turn_lm_masked_label).exp() lm_ppls.append(lm_ppl) turn_lm_masked_wrong_mask = turn_lm_masked_logit.argmax(dim=-1) != turn_lm_masked_label if turn_lm_masked_wrong_mask.any(): num_lm_correct_tokens = turn_lm_masked_wrong_mask.nonzero()[0,0] else: num_lm_correct_tokens = (~turn_lm_masked_wrong_mask).sum() lm_correctness.append(num_lm_correct_tokens / turn_lm_masked_label.numel()) ## 3.3. frame_diff (will be casted to time_diff in compute_metrics) if turn_stream_mask.any(): ## 3.3.1: reply before (at) turn_num_frames turn_score = turn_logit.softmax(dim=-1) turn_stream_masked_score = turn_score[turn_stream_mask] if frame_token_interval_threshold > 0: lower_threshold_mask = turn_stream_masked_score[:, frame_token_interval_id] < frame_token_interval_threshold turn_stream_masked_score[lower_threshold_mask] = 0 turn_stream_masked_pred_mask = turn_stream_masked_score.argmax(dim=-1) != frame_token_interval_id if turn_stream_masked_pred_mask.any(): frame_diff = turn_stream_mask.sum() - turn_stream_masked_pred_mask.nonzero()[0,0] - 1 else: ## 3.3.2: the most complex part,reply after turn_num_frames. we assume the 'assistant: ...' not exists turn_last_stream_idx = turn_stream_mask.nonzero()[-1,0] past_key_values_before_assistant = self.trim_past_key_values(past_key_values, 0, turn_start + turn_last_stream_idx + 1) if r == num_turns - 1: # no future frame. we assume the model should receive a signal when streaming ends (e.g. close button). frame_diff = zero else: next_turn_num_frames = (input_id[turn_starts[r+1]:turn_stops[r+1]] == v_placeholder_id).sum() // frame_num_tokens to_append_num_frames = min(next_turn_num_frames, turn_num_frames - 1) # avoid bias. current as center, two equal left/right side if to_append_num_frames == 0: frame_diff = zero else: to_append_frames = frames[past_num_frames+turn_num_frames:past_num_frames+turn_num_frames+to_append_num_frames] frame_placeholder = [v_placeholder_id] * frame_num_tokens if use_interval: frame_placeholder = [frame_token_interval_id] + frame_placeholder to_append_input_id = torch.tensor(frame_placeholder * to_append_num_frames, dtype=torch.long, device=device) to_append_logit = self.forward( input_ids=to_append_input_id[None], past_key_values=past_key_values_before_assistant, frames=to_append_frames, return_dict=True, use_cache=True ).logits[0] # we only use the last idx of each frame idxs = torch.arange(len(frame_placeholder)-1, len(to_append_input_id), len(frame_placeholder), device=device) to_append_score = to_append_logit[idxs].softmax(dim=-1) if frame_token_interval_threshold > 0: lower_threshold_mask = to_append_score[:, frame_token_interval_id] < frame_token_interval_threshold to_append_score[lower_threshold_mask] = 0 to_append_score_pred_mask = to_append_score.argmax(dim=-1) != frame_token_interval_id if to_append_score_pred_mask.any(): frame_diff = -(to_append_score_pred_mask.nonzero()[0,0] + 1) else: frame_diff = -to_append_num_frames frame_diffs.append(frame_diff.abs()) ## 2.6 fluency if turn_lm_mask.any() and turn_stream_mask.any(): num_learn_v_tokens = turn_stream_mask.sum() num_learn_valid_tokens = turn_lm_masked_label.numel() + num_learn_v_tokens if frame_diff == 0: fluency = (num_learn_v_tokens + num_lm_correct_tokens) / num_learn_valid_tokens elif frame_diff > 0: fluency = (num_learn_v_tokens - frame_diff) / num_learn_valid_tokens else: fluency = (num_learn_v_tokens - 1) / num_learn_valid_tokens fluencies.append(fluency) ## 2.7 next turn past_num_frames += turn_num_frames lm_ppl = torch.stack(lm_ppls).mean() if lm_ppls else one frame_diff = torch.stack(frame_diffs).float().mean() if frame_diffs else zero fluency = torch.stack(fluencies).float().mean() if fluencies else one lm_correctness = torch.stack(lm_correctness).float().mean() if lm_correctness else one return torch.stack([lm_ppl, frame_diff, fluency, lm_correctness]) def trim_past_key_values(self, past_key_values, start, stop): return [[past_keys[:,:,start:stop], past_values[:,:,start:stop]] for past_keys, past_values in past_key_values] def fast_greedy_generate(*, model: LiveMixin, inputs_embeds: torch.Tensor, past_key_values: Cache, eos_token_id: int, inplace_output_ids: torch.Tensor): for i in range(inplace_output_ids.size(1)): outputs = model(inputs_embeds=inputs_embeds, past_key_values=past_key_values, use_cache=True) past_key_values = outputs.past_key_values new_token_id = outputs.logits[:, -1:].argmax(dim=-1) inplace_output_ids[:, i] = new_token_id if new_token_id == eos_token_id: break inputs_embeds = model.get_input_embeddings()(new_token_id) return inplace_output_ids[:, :i+1], past_key_values def build_live( *, is_training: bool, config_class: type, model_class: type, llm_pretrained: str = None, finetune_modules: list[str] = None, lora_modules: str = None, lora_r: int = None, lora_alpha: int = None, set_vision_inside: bool = False, resume_from_checkpoint: str = '', attn_implementation: str = 'flash_attention_2', torch_dtype: str | torch.dtype = 'auto', **kwargs ): model = model_class.from_pretrained(llm_pretrained, config=config_class.from_pretrained(llm_pretrained, **kwargs), torch_dtype=torch_dtype, attn_implementation=attn_implementation) tokenizer = build_live_tokenizer_and_update_config(llm_pretrained, model.config) if is_training: lora_config = LoraConfig( r=lora_r, lora_alpha=lora_alpha, target_modules=lora_modules, lora_dropout=0.05, task_type="CAUSAL_LM", modules_to_save=finetune_modules, inference_mode=False, ) model = get_peft_model(model, lora_config) model.print_trainable_parameters() else: if resume_from_checkpoint and os.path.exists(resume_from_checkpoint): model = PeftModel.from_pretrained(model, resume_from_checkpoint, is_trainable=False) else: logger.warning(f'!!! Fail to load checkpoint: {resume_from_checkpoint}. Return a new initialized model.') if set_vision_inside: model.set_vision_inside() model.requires_grad_(False) return model, tokenizer