| 10 | |
| 11 | |
| 12 | class PromptAdapterManager: |
| 13 | def __init__(self): |
| 14 | self._uid_counter = 0 |
| 15 | self._uid_to_weights: Dict[str, torch.Tensor] = {} |
| 16 | |
| 17 | def load_from_ckpt( |
| 18 | self, model_dirs: List[str], model_config: "ModelConfig", uids: Optional[List[str]] = None |
| 19 | ): |
| 20 | if uids is None: |
| 21 | uids = [self._generate_uid() for _ in range(len(model_dirs))] |
| 22 | assert len(uids) == len(model_dirs) |
| 23 | |
| 24 | new_uids, new_model_dirs = [], [] |
| 25 | for uid, model_dir in zip(uids, model_dirs): |
| 26 | if uid in self._uid_to_weights: |
| 27 | continue |
| 28 | new_uids.append(uid) |
| 29 | new_model_dirs.append(model_dir) |
| 30 | |
| 31 | if len(new_uids) == 0: |
| 32 | return |
| 33 | |
| 34 | for uid, model_dir in zip(new_uids, new_model_dirs): |
| 35 | state_dict = load_state_dict(get_model_path(model_dir, "adapter_model")) |
| 36 | self._uid_to_weights[uid] = state_dict["prompt_embeddings"].to( |
| 37 | str_dtype_to_torch(model_config.dtype) |
| 38 | ) |
| 39 | |
| 40 | @property |
| 41 | def uid_to_weights(self): |
| 42 | return self._uid_to_weights |
| 43 | |
| 44 | def _generate_uid(self): |
| 45 | while str(self._uid_counter) in self._uid_to_weights: |
| 46 | self._uid_counter += 1 |
| 47 | uid = str(self._uid_counter) |
| 48 | self._uid_counter += 1 |
| 49 | return uid |