Update modeling_quiet.py

modeling_quiet.py

@@ -1424,6 +1424,46 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
         # 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,
+		streamer: Optional[TextStreamer] = None,
+		use_cache: Optional[bool] = None,
+		output_attentions: Optional[bool] = None,
+		output_hidden_states: Optional[bool] = None,
+		return_dict: Optional[bool] = None,
+		**kwargs,
+	):
+		# Call your custom infer function
+		logits = self.infer(
+			input_ids,
+			attention_mask,
+			position_ids,
+			past_key_values,
+			inputs_embeds,
+			use_cache,
+			output_attentions,
+			output_hidden_states,
+			return_dict,
+		)
+
+		# Sample the next token using the logits
+		next_token = torch.multinomial(torch.softmax(logits, dim=-1), num_samples=1)
+
+		# Append the generated token to the input sequence
+		input_ids = torch.cat([input_ids, next_token], dim=-1)
+
+		# Stream the generated token if a streamer is provided
+		if streamer is not None:
+			streamer.put(next_token)
+
+		return input_ids
 
     @add_start_docstrings_to_model_forward(QUIET_INPUTS_DOCSTRING)
     @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)