Update modeling_quiet.py

modeling_quiet.py

@@ -1182,7 +1182,6 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		self.gradient_accumulation_steps = 1
 		self.training_steps = 0
 		self.tokenizer = AutoTokenizer.from_pretrained("Crystalcareai/Quiet-Star-Custom")
-		self.streamer = TextStreamer(self.tokenizer)
 		self.start_token_id = None
 		self.end_token_id = None
 		self.rm_initialized = False
@@ -2116,6 +2115,15 @@ class QuietForCausalLM(QuietPreTrainedModel, GenerationMixin):
 		del start_embedding
 		del end_embedding
 		torch.cuda.empty_cache()
+
+		if streamer is not None:
+			streamer_kwargs = {
+				"generated_token_ids": input_ids,
+				"past_key_values": outputs.past_key_values,
+				"hidden_states": outputs.hidden_states,
+				"attentions": outputs.attentions,
+			}
+			streamer.put(**streamer_kwargs)
 		
 
 		return CausalLMOutputWithPast(