(
model: PreTrainedModel,
tokenizer: PreTrainedTokenizer,
messages: List[Dict[str, str]],
max_new_tokens: int = 512,
temperature: float = 0.666,
top_p: float = 0.90,
top_k: int = 27,
min_p: float = 0.03,
generator: torch.Generator = torch.Generator(device=device).manual_seed(1337)
)
| 123 | return samples[best_sample_idx] |
| 124 | |
| 125 | def entropy_decode( |
| 126 | model: PreTrainedModel, |
| 127 | tokenizer: PreTrainedTokenizer, |
| 128 | messages: List[Dict[str, str]], |
| 129 | max_new_tokens: int = 512, |
| 130 | temperature: float = 0.666, |
| 131 | top_p: float = 0.90, |
| 132 | top_k: int = 27, |
| 133 | min_p: float = 0.03, |
| 134 | generator: torch.Generator = torch.Generator(device=device).manual_seed(1337) |
| 135 | ) -> str: |
| 136 | model.to(device) |
| 137 | logging.info("Starting entropy decoding") |
| 138 | |
| 139 | if hasattr(tokenizer, 'chat_template') and tokenizer.chat_template: |
| 140 | input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) |
| 141 | else: |
| 142 | input_text = "\n".join([f"{msg['role']}: {msg['content']}" for msg in messages]) |
| 143 | input_text += "\nassistant:" |
| 144 | |
| 145 | input_ids = tokenizer.encode(input_text, return_tensors="pt").to(device) |
| 146 | attention_mask = torch.ones_like(input_ids).to(device) |
| 147 | |
| 148 | if tokenizer.pad_token_id is None: |
| 149 | tokenizer.pad_token_id = tokenizer.eos_token_id |
| 150 | |
| 151 | generated_tokens = [] |
| 152 | gen_tokens = input_ids |
| 153 | past_key_values = None |
| 154 | stop = torch.tensor([tokenizer.eos_token_id], device=device, dtype=torch.int32) |
| 155 | |
| 156 | for step in range(max_new_tokens): |
| 157 | logging.debug(f"Generation step: {step + 1}") |
| 158 | with torch.no_grad(): |
| 159 | outputs = model( |
| 160 | input_ids if past_key_values is None else input_ids[:, -1:], |
| 161 | attention_mask=attention_mask, |
| 162 | past_key_values=past_key_values, |
| 163 | use_cache=True, |
| 164 | output_attentions=True, |
| 165 | ) |
| 166 | |
| 167 | logits = outputs.logits[:, -1:, :] |
| 168 | attention_scores = outputs.attentions[-1] |
| 169 | past_key_values = outputs.past_key_values |
| 170 | |
| 171 | entropy, varentropy = calculate_varentropy_logsoftmax(logits) |
| 172 | attention_metrics = calculate_attention_metrics(attention_scores) |
| 173 | metrics = { |
| 174 | "logits_entropy": entropy, |
| 175 | "logits_varentropy": varentropy, |
| 176 | **attention_metrics |
| 177 | } |
| 178 | |
| 179 | logging.debug(f"Metrics: entropy={entropy.item():.3f}, varentropy={varentropy.item():.3f}") |
| 180 | |
| 181 | if entropy < 0.1 and varentropy < 0.1: |
| 182 | next_token = torch.argmax(logits[:, -1], dim=-1, keepdim=True).to(torch.int32) |
no test coverage detected