Archives/petals
2
0
Fork 0
mirror of https://github.com/bigscience-workshop/petals synced 2024-10-31 09:20:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

Update block.py

Browse Source
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
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
src/petals/bloom/block.py
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):