Update modeling_quiet.py

modeling_quiet.py

@@ -23,7 +23,7 @@ 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
@@ -32,7 +32,7 @@ from torch import nn
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 from transformers.generation.utils import GenerationMixin
 from transformers.generation.stopping_criteria import StoppingCriteriaList, validate_stopping_criteria
-from transformers import TextStreamer
+from transformers import TextStreamer, AutoTokenizer
 
 from transformers.activations import ACT2FN
 from transformers.cache_utils import Cache, DynamicCache
@@ -774,7 +774,7 @@ class QuietSdpaAttention(QuietAttention):
 				raise ValueError(
 					f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
 				)
-
+			attention_mask = attention_mask.to(query_states.dtype)
 		# SDPA with memory-efficient backend is currently (torch==2.1.2) bugged with non-contiguous inputs with custom attn_mask,
 		# Reference: https://github.com/pytorch/pytorch/issues/112577.
 		if query_states.device.type == "cuda" and attention_mask is not None:
@@ -1182,7 +1182,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		self.n_tokens_print = 1
 		self.gradient_accumulation_steps = 1
 		self.training_steps = 0
-		self.tokenizer = None
+		self.tokenizer = AutoTokenizer.from_pretrained("Crystalcareai/Quiet-Star-Custom")
 		self.start_token_id = None
 		self.end_token_id = None
 		self.rm_initialized = False
@@ -1238,7 +1238,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		self.embedding_scale = 1e2
 		self.temperature = nn.Parameter(torch.ones(1))
 		self.max_temperature = config.max_temperature
-		self.complexity_factor = config.complexity_factor
 		self.reinforce_temperature = 3
 		self.base_loss_beta = 1
 		self.thinking_usefulness_head = nn.Linear(self.model.config.hidden_size, 1)
@@ -1425,67 +1424,9 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		logits = self.lm_head(mixed_hidden_states)
 		return logits
 
-	@torch.no_grad()
-	def generate(
-		self,
-		input_ids: Optional[torch.LongTensor] = None,
-		max_length: Optional[int] = None,
-		min_length: Optional[int] = None,
-		do_sample: Optional[bool] = None,
-		early_stopping: Optional[bool] = None,
-		num_beams: Optional[int] = None,
-		temperature: Optional[float] = None,
-		top_k: Optional[int] = None,
-		top_p: Optional[float] = None,
-		repetition_penalty: Optional[float] = None,
-		bad_words_ids: Optional[Iterable[int]] = None,
-		bos_token_id: Optional[int] = None,
-		pad_token_id: Optional[int] = None,
-		eos_token_id: Optional[int] = None,
-		length_penalty: Optional[float] = None,
-		no_repeat_ngram_size: Optional[int] = None,
-		encoder_no_repeat_ngram_size: Optional[int] = None,
-		num_return_sequences: Optional[int] = None,
-		max_time: Optional[float] = None,
-		max_new_tokens: Optional[int] = None,
-		decoder_start_token_id: Optional[int] = None,
-		use_cache: Optional[bool] = None,
-		num_beam_groups: Optional[int] = None,
-		diversity_penalty: Optional[float] = None,
-		prefix_allowed_tokens_fn: Optional[Callable[[int, torch.Tensor], List[int]]] = None,
-		output_attentions: Optional[bool] = None,
-		output_hidden_states: Optional[bool] = None,
-		output_scores: Optional[bool] = None,
-		return_dict_in_generate: Optional[bool] = None,
-		forced_bos_token_id: Optional[int] = None,
-		forced_eos_token_id: Optional[int] = None,
-		remove_invalid_values: Optional[bool] = None,
-		synced_gpus: Optional[bool] = False,
-		stopping_criteria: Optional[StoppingCriteriaList] = None,
-		**model_kwargs,
-	):
-		# # Validate stopping criteria
-		# stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)
-
-		# Set default values for attention_mask and max_new_tokens
-		if "attention_mask" not in model_kwargs:
-			attention_mask = torch.where(input_ids != self.tokenizer.pad_token_id, 1, 0).to(input_ids.device)
-			model_kwargs["attention_mask"] = attention_mask
-		if max_new_tokens is None:
-			max_new_tokens = 512
-
-		streamer = TextStreamer(self.tokenizer, skip_prompt=False, skip_special_tokens=True)
-
-		# Call the custom generate function
-		output_ids, _ = custom_generate(
-			self,
-			input_ids=input_ids,
-			streamer=streamer,
-			max_new_tokens=max_new_tokens,
-			**model_kwargs,
-		)
-
-		return output_ids
+	def generate(self, input_ids, attention_mask=None, max_length=None, temperature=1.0, **kwargs):
+		from .generate import generate
+		return generate(self, input_ids=input_ids, attention_mask=attention_mask, max_length=max_length, temperature=temperature, **kwargs)
 
 	@add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
 	@replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)
@@ -1675,8 +1616,7 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		sample_probs_history = []
 		action_loglikelihoods_list = []
 
-		complexity_scores = self.compute_complexity_scores(input_ids, attention_mask)
-		temperature = self.temperature * complexity_scores.unsqueeze(-1)
+		temperature = self.temperature
 
 		if self.use_end_thought_token or self.use_start_thought_token:
 			if not self.use_reparam_for_thought_embeddings:
@@ -1734,10 +1674,10 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 					base_attention_mask = base_attention_mask.view(1, 1, seq_len, seq_len)
 					base_attention_mask = base_attention_mask.repeat(input_ids.shape[0], 1, 1, 1)
 					attention_mask = base_attention_mask
-					breakpoint()
+					# breakpoint()
 				elif attention_mask.dim() == 2:
 					if seq_len + past_key_values_length != attention_mask.shape[-1]:
-						breakpoint()
+						# breakpoint()
 						attention_mask = torch.cat(
 							[torch.ones((attention_mask.shape[0], past_key_values_length), dtype=attention_mask.dtype, device=attention_mask.device), attention_mask],
 							dim=-1
@@ -2190,33 +2130,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 			attentions=outputs.attentions,
 		)
 		
-	
-		
-	def compute_complexity_scores(self, input_ids, attention_mask):
-		# Compute complexity scores based on input sequence characteristics
-		# Example: Normalize sequence lengths and consider the presence of rare tokens
-		seq_lengths = torch.sum(attention_mask, dim=-1)
-		max_length = torch.max(seq_lengths)
-		length_scores = seq_lengths / max_length
-
-		# Compute the proportion of rare tokens in each sequence
-		rare_token_ids = self.get_rare_token_ids()
-		rare_token_mask = torch.isin(input_ids, rare_token_ids)
-		rare_token_counts = torch.sum(rare_token_mask, dim=-1)
-		rare_token_scores = rare_token_counts / seq_lengths
-
-		# Combine length scores and rare token scores
-		complexity_scores = self.complexity_factor * length_scores + (1 - self.complexity_factor) * rare_token_scores
-		return complexity_scores
-
-	def get_rare_token_ids(self):
-		# Get the IDs of rare tokens based on a predefined frequency threshold
-		frequency_threshold = 1e-4
-		token_counts = torch.bincount(self.model.embed_tokens.weight.argmax(dim=-1))
-		total_tokens = torch.sum(token_counts)
-		rare_token_mask = token_counts / total_tokens < frequency_threshold
-		rare_token_ids = torch.nonzero(rare_token_mask).squeeze(-1)
-		return rare_token_ids
 
 
 	def prepare_inputs_for_generation(