Update custom_generate/generate.py

custom_generate/generate.py

@@ -30,10 +30,14 @@ import torch.nn as nn
 from transformers.modeling_utils import  PreTrainedModel
 
 
-
-from ..functions_2_patch import _validate_model_kwargs, llama_atten_forward
-from ..monkey_patching_utils import monkey_patching
-from ..sep_cache_utils import SepCache
+try:
+    from .functions_2_patch import _validate_model_kwargs, llama_atten_forward
+    from .monkey_patching_utils import monkey_patching
+    from .sep_cache_utils import SepCache
+except ImportError as e:
+    from ..functions_2_patch import _validate_model_kwargs, llama_atten_forward
+    from ..monkey_patching_utils import monkey_patching
+    from ..sep_cache_utils import SepCache
 
 
 UNSUPPORTED_GENERATION_ARGS = [