(models_plus: list[ModelPlus])
| 899 | |
| 900 | |
| 901 | def merge_multifile_models(models_plus: list[ModelPlus]) -> ModelPlus: |
| 902 | formats = set(mp.format for mp in models_plus) |
| 903 | assert len(formats) == 1, "different formats?" |
| 904 | format = formats.pop() |
| 905 | paths = [path for mp in models_plus for path in mp.paths] |
| 906 | # Use the first non-None vocab, if any. |
| 907 | try: |
| 908 | vocab = next(mp.vocab for mp in models_plus if mp.vocab is not None) |
| 909 | except StopIteration: |
| 910 | vocab = None |
| 911 | |
| 912 | if any("model.embed_tokens.weight" in mp.model for mp in models_plus): |
| 913 | # Transformers models put different tensors in different files, but |
| 914 | # don't split individual tensors between files. |
| 915 | model: LazyModel = {} |
| 916 | for mp in models_plus: |
| 917 | model.update(mp.model) |
| 918 | else: |
| 919 | model = merge_sharded([mp.model for mp in models_plus]) |
| 920 | |
| 921 | return ModelPlus(model, paths, format, vocab) # pytype: disable=wrong-arg-types |
| 922 | |
| 923 | |
| 924 | def permute_lazy(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int) -> LazyTensor: |
no test coverage detected