MCPcopy Index your code
hub / github.com/NVIDIA/TensorRT-LLM / __init__

Method __init__

tensorrt_llm/lora_manager.py:251–287  ·  view source on GitHub ↗
(self, lora_dirs: List[str])

Source from the content-addressed store, hash-verified

249
250class HfLoraLoader:
251 def __init__(self, lora_dirs: List[str]):
252 self.lora_target_modules = []
253 self.is_valid = False
254 self.lm_head = None
255 self.embed_tokens = None
256 self.vocab_size = 0
257
258 if len(lora_dirs) == 0:
259 return
260
261 for lora_dir in lora_dirs:
262 model_path = get_model_path(lora_dir, "adapter_model")
263 if model_path is None:
264 raise ValueError(f"adapter_model file does not exist in {lora_dir}")
265 config_file = Path(f"{lora_dir}/adapter_config.json")
266 if not config_file.exists():
267 raise ValueError(f"{config_file} does not exist")
268 if not config_file.is_file():
269 raise ValueError(f"{config_file} is not a file")
270 self.is_valid = True
271
272 lora_dir = lora_dirs[0]
273 with open(f"{lora_dir}/adapter_config.json") as f:
274 adapter_config = json.load(f)
275
276 model_path = get_model_path(lora_dir, "adapter_model")
277 if model_path is None:
278 raise ValueError(f"adapter_model file does not exist in {lora_dir}")
279 lora_weight = load_state_dict(model_path)
280 self.lora_weight = lora_weight
281 if adapter_config.get("modules_to_save") is not None:
282 if "lm_head" in adapter_config["modules_to_save"]:
283 self.lm_head = lora_weight["base_model.model.lm_head.weight"]
284 self.vocab_size = self.lm_head.shape[0]
285
286 if "embed_tokens" in adapter_config["modules_to_save"]:
287 self.embed_tokens = lora_weight["base_model.model.model.embed_tokens.weight"]
288
289 def get_target_modules(self, trtllm_modules_to_hf_modules):
290 hf_modules_to_trtllm_modules = invert_module_mapping(trtllm_modules_to_hf_modules)

Callers

nothing calls this directly

Calls 4

load_state_dictFunction · 0.85
get_model_pathFunction · 0.50
loadMethod · 0.45
getMethod · 0.45

Tested by

no test coverage detected