(models_plus: list[ModelPlus])
| 713 | |
| 714 | |
| 715 | def merge_multifile_models(models_plus: list[ModelPlus]) -> ModelPlus: |
| 716 | formats = set(mp.format for mp in models_plus) |
| 717 | assert len(formats) == 1, "different formats?" |
| 718 | format = formats.pop() |
| 719 | paths = [path for mp in models_plus for path in mp.paths] |
| 720 | # Use the first non-None vocab, if any. |
| 721 | try: |
| 722 | vocab = next(mp.vocab for mp in models_plus if mp.vocab is not None) |
| 723 | except StopIteration: |
| 724 | vocab = None |
| 725 | |
| 726 | if any("model.embed_tokens.weight" in mp.model for mp in models_plus): |
| 727 | # Transformers models put different tensors in different files, but |
| 728 | # don't split individual tensors between files. |
| 729 | model: LazyModel = {} |
| 730 | for mp in models_plus: |
| 731 | model.update(mp.model) |
| 732 | else: |
| 733 | model = merge_sharded([mp.model for mp in models_plus]) |
| 734 | |
| 735 | return ModelPlus(model, paths, format, vocab) # pytype: disable=wrong-arg-types |
| 736 | |
| 737 | |
| 738 | def permute_lazy(lazy_tensor: LazyTensor, n_head: int, n_head_kv: int) -> LazyTensor: |
no test coverage detected