A light wrapper around **two** causal‑LMs that fuses their logits: logits_joint = logaddexp(logits_base + log(1‑λ), logits_knn + log(λ)) Greedy decoding chooses argmax over `logits_joint`.
| 28 | attentions: Optional[Tuple[torch.FloatTensor, ...]] = None |
| 29 | |
| 30 | class MemoryDecoder(PreTrainedModel, GenerationMixin): |
| 31 | """ |
| 32 | A light wrapper around **two** causal‑LMs that fuses their logits: |
| 33 | |
| 34 | logits_joint = logaddexp(logits_base + log(1‑λ), |
| 35 | logits_knn + log(λ)) |
| 36 | |
| 37 | Greedy decoding chooses argmax over `logits_joint`. |
| 38 | """ |
| 39 | |
| 40 | def __init__( |
| 41 | self, |
| 42 | base_lm, |
| 43 | knn_generator, |
| 44 | lmbda: float = 0.25, |
| 45 | knn_temp: float = 1.0, |
| 46 | ): |
| 47 | super().__init__(base_lm.config) |
| 48 | |
| 49 | self.base_lm = base_lm |
| 50 | self.knn_generator = knn_generator |
| 51 | self.lmbda = float(lmbda) |
| 52 | self.knn_temp = float(knn_temp) |
| 53 | |
| 54 | # ------------------------------------------------------------------ # |
| 55 | # 1. forward() |
| 56 | # ------------------------------------------------------------------ # |
| 57 | def forward( |
| 58 | self, |
| 59 | input_ids: torch.LongTensor, |
| 60 | attention_mask: Optional[torch.LongTensor] = None, |
| 61 | past_key_values: Optional[Tuple] = None, |
| 62 | knn_past_key_values: Optional[Tuple] = None, |
| 63 | use_cache: bool = True, |
| 64 | **kwargs, |
| 65 | ): |
| 66 | """ |
| 67 | Forward pass that returns **fused log‑probs** as logits. |
| 68 | We keep separate caches for each sub‑model. |
| 69 | """ |
| 70 | base_outputs = self.base_lm( |
| 71 | input_ids=input_ids, |
| 72 | attention_mask=attention_mask, |
| 73 | past_key_values=past_key_values, |
| 74 | use_cache=use_cache, |
| 75 | **kwargs, |
| 76 | ) |
| 77 | knn_outputs = self.knn_generator( |
| 78 | input_ids=input_ids, |
| 79 | attention_mask=attention_mask, |
| 80 | past_key_values=knn_past_key_values, |
| 81 | use_cache=use_cache, |
| 82 | **kwargs, |
| 83 | ) |
| 84 | |
| 85 | # Temperature on k‑NN logits only |
| 86 | logits_base = base_outputs.logits # (B, T, V) |
| 87 | logits_knn = knn_outputs.logits |