Update modeling_quiet.py

modeling_quiet.py

@@ -23,16 +23,16 @@ import math
 import pdb
 import warnings
 from collections import defaultdict
-from typing import List, Optional, Tuple, Union, Iterable, Callable
+from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn.functional as F
 import torch.utils.checkpoint
 from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
-from transformers.generation.utils import GenerationMixin, GenerationConfig
+from transformers.generation.utils import GenerationMixin
 from transformers.generation.stopping_criteria import StoppingCriteriaList, validate_stopping_criteria
-from transformers import TextStreamer, AutoTokenizer
+from transformers import TextStreamer
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
@@ -143,6 +143,7 @@ class QuietRMSNorm(nn.Module):
 		self.weight = nn.Parameter(torch.ones(hidden_size))
 		self.variance_epsilon = eps
 
+
 	def forward(self, hidden_states):
 		input_dtype = hidden_states.dtype
 		hidden_states = hidden_states.to(torch.float32)
@@ -150,6 +151,7 @@ class QuietRMSNorm(nn.Module):
 		hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
 		return hidden_states.to(input_dtype) * self.weight.to(hidden_states.device)
 
+
 # Copied from transformers.models.llama.modeling_llama.LlamaRotaryEmbedding with Llama->Quiet
 class QuietRotaryEmbedding(nn.Module):
 	def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
@@ -235,8 +237,7 @@ class QuietMLP(nn.Module):
 		self.act_fn = ACT2FN[config.hidden_act]
 
 	def forward(self, x):
-		hidden_states = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
-		return hidden_states
+		return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
 
 
 # Copied from transformers.models.llama.modeling_llama.repeat_kv
@@ -848,7 +849,7 @@ class QuietDecoderLayer(nn.Module):
 		residual = hidden_states
 
 		hidden_states = self.input_layernorm(hidden_states)
-		
+
 		# Self Attention
 		hidden_states, self_attn_weights, present_key_value = self.self_attn(
 			hidden_states=hidden_states,
@@ -1022,8 +1023,6 @@ class QuietModel(QuietPreTrainedModel):
 		output_hidden_states: Optional[bool] = None,
 		return_dict: Optional[bool] = None,
 	) -> Union[Tuple, BaseModelOutputWithPast]:
-		
-		# print("Hidden states shape after embedding:", inputs_embeds.shape)
 		output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
 		output_hidden_states = (
 			output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
@@ -1072,27 +1071,32 @@ class QuietModel(QuietPreTrainedModel):
 		if self._attn_implementation == "flash_attention_2":
 			# 2d mask is passed through the layers
 			attention_mask = attention_mask if (attention_mask is not None and 0 in attention_mask) else None
-		elif self._attn_implementation == "sdpa" and not output_attentions and attention_mask is not None and attention_mask.dim() == 2 and False:
+		elif self._attn_implementation == "sdpa" and not output_attentions and attention_mask.dim() == 2 and False:
 			# output_attentions=True can not be supported when using SDPA, and we fall back on
 			# the manual implementation that requires a 4D causal mask in all cases.
 			attention_mask = _prepare_4d_causal_attention_mask_for_sdpa(
-				attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length,
+				attention_mask,
+				(batch_size, seq_length),
+				inputs_embeds,
+				past_key_values_length,
 			)
 		elif attention_mask is None or attention_mask.dim() == 2:
 			# 4d mask is passed through the layers
 			attention_mask = _prepare_4d_causal_attention_mask(
-				attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length,
+				attention_mask,
+				(batch_size, seq_length),
+				inputs_embeds,
+				past_key_values_length,
 				sliding_window=self.config.sliding_window,
 			)
 
-
 		hidden_states = inputs_embeds
 
 		# decoder layers
 		all_hidden_states = () if output_hidden_states else None
 		all_self_attns = () if output_attentions else None
 		next_decoder_cache = None
-		# print("Hidden states shape before decoder layers:", hidden_states.shape)
+
 		for decoder_layer in self.layers:
 			if output_hidden_states:
 				all_hidden_states += (hidden_states,)
@@ -1116,15 +1120,15 @@ class QuietModel(QuietPreTrainedModel):
 					output_attentions=output_attentions,
 					use_cache=use_cache,
 				)
+
 			hidden_states = layer_outputs[0]
-			# print(f"Hidden states shape after decoder layer {decoder_layer}:", hidden_states.shape)
-			# print("Hidden states shape after decoder layers:", hidden_states.shape)
 
 			if use_cache:
 				next_decoder_cache = layer_outputs[2 if output_attentions else 1]
 
 			if output_attentions:
 				all_self_attns += (layer_outputs[1],)
+
 		hidden_states = self.norm(hidden_states)
 
 		# add hidden states from the last decoder layer
@@ -1155,7 +1159,7 @@ def loss_mean(x):
 class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 	_tied_weights_keys = ["lm_head.weight"]
 
-	def __init__(self, config,tokenizer=None):
+	def __init__(self, config):
 		super().__init__(config)
 		self.model = QuietModel(config)
 		self.vocab_size = config.vocab_size
@@ -1178,7 +1182,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		self.n_tokens_print = 1
 		self.gradient_accumulation_steps = 1
 		self.training_steps = 0
-		self.tokenizer = tokenizer
+		self.tokenizer = None
 		self.start_token_id = None
 		self.end_token_id = None
 		self.rm_initialized = False
@@ -1306,14 +1310,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 				nn.init.constant_(module.bias, 0)
 		elif isinstance(module, nn.Embedding):
 			nn.init.xavier_uniform_(module.weight)
-			
-	@classmethod
-	def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
-		tokenizer = AutoTokenizer.from_pretrained(pretrained_model_name_or_path)
-		model = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
-		model.tokenizer = tokenizer
-		return model
-
 
 	@torch.no_grad()
 	def infer(
@@ -1347,10 +1343,13 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		continuation_length = self.n_ahead - 2
 		new_key_values = past_key_values
 		
+		# Initialize generated_ids with input_ids
+		generated_ids = input_ids.clone()
+
 		start_time = time.time()
 		for continuation_idx in range(continuation_length):
 			outputs = self.model(
-				input_ids=input_ids if continuation_idx == 0 else next_token_id.unsqueeze(-1).to(input_ids.device),
+				input_ids=generated_ids if continuation_idx == 0 else next_token_id.unsqueeze(-1).to(generated_ids.device),
 				attention_mask=attention_mask,
 				position_ids=position_ids,
 				past_key_values=new_key_values,
@@ -1371,86 +1370,33 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 			next_token_logits = F.gumbel_softmax(logits, tau=self.gumbel_temperature, hard=True, dim=-1)
 			next_token_id = torch.argmax(next_token_logits, dim=-1)
 
-			# Append the generated token to the input sequence
-			input_ids = torch.cat([input_ids, next_token_id.unsqueeze(-1).to(input_ids.device)], dim=-1)
-			seq_len += 1
+			# Append the generated token to the generated_ids
+			generated_ids = torch.cat([generated_ids, next_token_id.unsqueeze(-1).to(generated_ids.device)], dim=-1)
 
-			# Update the attention mask
-			if attention_mask is not None:
-				attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1)
-
-		# Append the end thought token to the input sequence
+		# Append the end thought token to the generated_ids
 		end_thought_token_id = self.tokenizer.convert_tokens_to_ids("<|endthought|>")
-		input_ids = torch.cat([input_ids, torch.tensor([[end_thought_token_id]] * batch_size).to(input_ids.device)], dim=-1)
-		seq_len += 1
+		generated_ids = torch.cat([generated_ids, torch.tensor([[end_thought_token_id]] * batch_size).to(generated_ids.device)], dim=-1)
 
-		# Update the attention mask
-		if attention_mask is not None:
-			attention_mask = torch.cat([attention_mask, torch.ones((batch_size, 1)).to(attention_mask.device)], dim=-1)
-
-		# Get the hidden states before and after the thought
-		outputs_before = self.model(
-			input_ids=original_input_ids,
-			attention_mask=original_attention_mask,
-			position_ids=position_ids,
-			past_key_values=past_key_values,
-			inputs_embeds=inputs_embeds,
-			use_cache=use_cache,
-			output_attentions=output_attentions,
-			output_hidden_states=output_hidden_states,
-			return_dict=return_dict,
-		)
-		hidden_states_before = outputs_before[0][:, -1:, :]
-
-		outputs_after = self.model(
-			input_ids=torch.cat([next_token_id.unsqueeze(-1).to(input_ids.device), torch.tensor([[end_thought_token_id]] * batch_size).to(input_ids.device)], dim=-1),
-			attention_mask=attention_mask,
-			position_ids=position_ids,
-			past_key_values=new_key_values,
-			inputs_embeds=inputs_embeds,
-			use_cache=use_cache,
-			output_attentions=output_attentions,
-			output_hidden_states=output_hidden_states,
-			return_dict=return_dict,
+		return generated_ids
+	
+	
+	@torch.no_grad()
+	def generate(self, *args, **kwargs):
+		# Call the infer method to generate the token ids
+		generated_ids = self.infer(
+			input_ids=kwargs.pop("input_ids", None),
+			attention_mask=kwargs.pop("attention_mask", None),
+			position_ids=kwargs.pop("position_ids", None),
+			past_key_values=kwargs.pop("past_key_values", None),
+			inputs_embeds=kwargs.pop("inputs_embeds", None),
+			use_cache=kwargs.pop("use_cache", None),
+			output_attentions=kwargs.pop("output_attentions", None),
+			output_hidden_states=kwargs.pop("output_hidden_states", None),
+			return_dict=kwargs.pop("return_dict", None),
 		)
-		hidden_states_after = outputs_after[0][:, -1:, :]
-
-		# Apply the talk head to get the mixing weight
-		mixing_weight = self.talk_head[0](torch.cat([hidden_states_before, hidden_states_after], dim=-1))
 
-		# Apply the mixing weight to the hidden states
-		mixed_hidden_states = (1 - mixing_weight) * hidden_states_before + mixing_weight * hidden_states_after
+		return generated_ids
 
-		# Apply the language model head to get the final logits
-		logits = self.lm_head(mixed_hidden_states)
-		
-		return logits
-
-	@torch.no_grad()
-	def generate(
-		self,
-		input_ids: torch.LongTensor,
-		attention_mask: Optional[torch.Tensor] = None,
-		position_ids: Optional[torch.LongTensor] = None,
-		past_key_values: Optional[List[torch.FloatTensor]] = None,
-		inputs_embeds: Optional[torch.FloatTensor] = None,
-		use_cache: Optional[bool] = None,
-		output_attentions: Optional[bool] = None,
-		output_hidden_states: Optional[bool] = None,
-		return_dict_in_generate: Optional[bool] = None,
-		**model_kwargs,
-	) -> Union[BaseModelOutputWithPast, torch.LongTensor]:
-		return self.infer(
-			input_ids=input_ids,
-			attention_mask=attention_mask,
-			position_ids=position_ids,
-			past_key_values=past_key_values,
-			inputs_embeds=inputs_embeds,
-			use_cache=use_cache,
-			output_attentions=output_attentions,
-			output_hidden_states=output_hidden_states,
-			return_dict=return_dict_in_generate,
-		)
 	@add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
 	@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
 	def forward(
@@ -1641,7 +1587,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 
 		complexity_scores = self.compute_complexity_scores(input_ids, attention_mask)
 		temperature = self.temperature * complexity_scores.unsqueeze(-1)
-		# pdb.set_trace()
+
 		if self.use_end_thought_token or self.use_start_thought_token:
 			if not self.use_reparam_for_thought_embeddings:
 				start_embedding = self.start_embedding[0].unsqueeze(0) * self.embedding_scale * temperature
@@ -1671,10 +1617,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 				position_ids = position_ids.unsqueeze(0).view(-1, seq_len)
 			else:
 				position_ids = position_ids.view(-1, seq_len).long()
-				
-				
-			# print("Input IDs shape:", input_ids.shape)
-			# print("Inputs embeds shape before embedding:", inputs_embeds.shape if inputs_embeds is not None else None)
+
 			if inputs_embeds is None:
 				contains_start = self.use_start_thought_token and (input_ids == self.start_token_id).any()
 				contains_end = self.use_end_thought_token and (input_ids == self.end_token_id).any()
@@ -1694,7 +1637,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 				else:
 					with torch.set_grad_enabled(not self.train_only_thinking_embedding):
 						inputs_embeds = self.model.embed_tokens(input_ids)
-						# print("Inputs embeds shape after embedding:", inputs_embeds.shape)
 			
 			if self.n_ahead != 1 or self.n_ahead_talk != 1 or self.comparison_mode:
 				if attention_mask is None: