Fix sneaky problem caused by caching

app.py

@@ -233,6 +233,11 @@ def generate(model, inputs, metric, window_len, max_new_tokens, **kwargs):
     model_kwargs = dict(use_cache=True)
     for i in range(max_steps):
         pbar.progress(i / max_steps, f"{i}/{max_steps}")
+
+        if input_ids.shape[1] == window_len:
+            model_kwargs.update(use_cache=False)
+            if "past_key_values" in model_kwargs:
+                del model_kwargs["past_key_values"]
         model_inputs = model.prepare_inputs_for_generation(input_ids, **model_kwargs)
         model_outputs = model(**model_inputs)
         model_kwargs = model._update_model_kwargs_for_generation(model_outputs, model_kwargs, is_encoder_decoder=False)
@@ -250,7 +255,6 @@ def generate(model, inputs, metric, window_len, max_new_tokens, **kwargs):
         input_ids = torch.cat([input_ids, torch.tensor([[next_token]])], dim=1)
         if input_ids.shape[1] > window_len:
             input_ids = input_ids[:, 1:]
-            model_kwargs.update(use_cache=False, past_key_values=None)
         if logprobs_window.shape[0] == window_len:
             logprobs.append(
                 logprobs_window[torch.arange(window_len), input_ids.squeeze(0)]