From adec684e14a5b2cf48a09c1cb2e372d1c36adffa Mon Sep 17 00:00:00 2001
From: "Mr.Fire" <fire15@126.com>
Date: Tue, 26 Mar 2024 15:14:01 +0800
Subject: [PATCH] Update build_model.py

change model_name to arg so we can input local dir like "/model/xxx/Mix..." to load without network.
---
 src/build_model.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/src/build_model.py b/src/build_model.py
index 12bfa0c..bebd73c 100644
--- a/src/build_model.py
+++ b/src/build_model.py
@@ -168,9 +168,9 @@ def build_model(
     quant_config: QuantConfig,
     offload_config: OffloadConfig,
     state_path: str,
+    model_name: "mistralai/Mixtral-8x7B-Instruct-v0.1",
 ):
-    model_name = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-
+    
     state_dict_00 = load_00_expert_state_dict(state_path, device)
 
     def _make_module():