m97j commited on
Commit
4e65de4
·
1 Parent(s): 3e43d9a

Update initializer.py to use explicit Hub filenames

Browse files
Files changed (1) hide show
  1. models/initializer.py +5 -5
models/initializer.py CHANGED
@@ -29,14 +29,14 @@ def download_llm() -> str:
29
  """
30
  local_model_path = hf_hub_download(
31
  repo_id=config.HF_MODEL_HUB,
32
- filename=os.path.basename(config.LLM_LOCAL_PATH),
33
  token=config.HF_TOKEN,
34
  local_dir=config.LLM_LOCAL_DIR
35
  )
36
 
37
  hf_hub_download(
38
  repo_id=config.HF_MODEL_HUB,
39
- filename=os.path.basename(config.CONFIG_LOCAL_PATH),
40
  token=config.HF_TOKEN,
41
  local_dir=config.CONFIG_LOCAL_DIR
42
  )
@@ -49,9 +49,9 @@ def download_reranker() -> str:
49
  """
50
  return hf_hub_download(
51
  repo_id=config.HF_MODEL_HUB,
52
- filename=os.path.basename(config.RERANKER_LOCAL_PATH),
53
  token=config.HF_TOKEN,
54
- local_dir=os.path.dirname(config.RERANKER_LOCAL_PATH)
55
  )
56
 
57
  def load_llm(local_model_path: str) -> CustomModel:
@@ -68,7 +68,7 @@ def load_llm(local_model_path: str) -> CustomModel:
68
  )
69
  except Exception:
70
  # Fallback: manual load if file naming is non-standard (e.g., .pth)
71
- _config = AutoConfig.from_pretrained(config.LOCAL_CONFIG_DIR)
72
  model = CustomModel(_config)
73
  state_dict = torch.load(local_model_path, map_location="cpu")
74
  model.load_state_dict(state_dict)
 
29
  """
30
  local_model_path = hf_hub_download(
31
  repo_id=config.HF_MODEL_HUB,
32
+ filename=config.HF_LLM_FILENAME,
33
  token=config.HF_TOKEN,
34
  local_dir=config.LLM_LOCAL_DIR
35
  )
36
 
37
  hf_hub_download(
38
  repo_id=config.HF_MODEL_HUB,
39
+ filename=config.HF_CONFIG_FILENAME,
40
  token=config.HF_TOKEN,
41
  local_dir=config.CONFIG_LOCAL_DIR
42
  )
 
49
  """
50
  return hf_hub_download(
51
  repo_id=config.HF_MODEL_HUB,
52
+ filename=config.HF_RERANKER_FILENAME,
53
  token=config.HF_TOKEN,
54
+ local_dir=config.RERANKER_LOCAL_DIR
55
  )
56
 
57
  def load_llm(local_model_path: str) -> CustomModel:
 
68
  )
69
  except Exception:
70
  # Fallback: manual load if file naming is non-standard (e.g., .pth)
71
+ _config = AutoConfig.from_pretrained(config.CONFIG_LOCAL_DIR)
72
  model = CustomModel(_config)
73
  state_dict = torch.load(local_model_path, map_location="cpu")
74
  model.load_state_dict(state_dict)