Update block.py

This commit is contained in:
justheuristic 2023-03-28 23:51:58 +03:00 committed by GitHub
parent 8bb8e86a6f
commit 063108e905
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -5,13 +5,15 @@ See commit history for authorship.
"""
import os
from typing import Optional, Tuple
from packaging import version
import torch.nn.quantized.dynamic.modules.linear
import transformers
from transformers.models.bloom.modeling_bloom import BloomBlock, _expand_mask, _make_causal_mask, build_alibi_tensor
if not os.getenv("PETALS_IGNORE_DEPENDENCY_VERSION"):
assert transformers.__version__.startswith("4.26."), "Please install transformers 4.26.1"
assert version.parse("4.26.0") < version.parse(transformers.__version__) < version.parse("5.0.0"), \
"Please install transformers >=4.26.0,<5.0.0"
class WrappedBloomBlock(BloomBlock):