Use Longformer

modeling_cocom.py

@@ -1,4 +1,4 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, PreTrainedModel, PretrainedConfig, AutoModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig, PreTrainedModel, PretrainedConfig, AutoModel,LongformerForCausalLM, LongformerTokenizer
 import torch
 import math 
 from peft import get_peft_model, LoraConfig, TaskType
@@ -263,7 +263,7 @@ class COCOM(PreTrainedModel):
             attention_mask=dec_attention_mask.to(device),
             do_sample=False,
             top_p=None,
-            max_new_tokens=max_new_tokens
+            max_new_tokens=min(max_new_tokens, 4096)
             )
         decoded = self.decoder_tokenizer.batch_decode(output_ids, skip_special_tokens=True)
         return decoded