Upload indigo/model.py with huggingface_hub

indigo/model.py

@@ -166,7 +166,12 @@ class GPT(nn.Module):
             tokens = torch.cat((tokens, next_id), dim=1)
             if i == max_new_tokens - 1:
                 break
-            logits, _, caches = self(next_id, caches=caches, return_caches=True)
+            if caches is not None and caches[0][0].size(2) >= self.config.block_size:
+                caches = None
+                idx_in = tokens[:, -self.config.block_size:]
+            else:
+                idx_in = next_id
+            logits, _, caches = self(idx_in, caches=caches, return_caches=True)
         return tokens
 
     def num_params(self):