davanstrien HF Staff commited on
Commit
734af3d
·
1 Parent(s): 90d6ec0

load weights manually

Browse files
Files changed (1) hide show
  1. generate_summaries_uv.py +13 -5
generate_summaries_uv.py CHANGED
@@ -25,7 +25,7 @@ os.environ["VLLM_ATTENTION_BACKEND"] = "FLASHINFER"
25
 
26
  import polars as pl
27
  from datasets import Dataset, load_dataset
28
- from huggingface_hub import login, dataset_info
29
  from tqdm.auto import tqdm
30
  from transformers import AutoTokenizer
31
  from vllm import LLM, SamplingParams
@@ -106,10 +106,18 @@ def generate_summaries(
106
  input_dataset_id, card_type, min_likes, min_downloads
107
  )
108
 
109
- # Initialize model and tokenizer
110
- logger.info(f"Initializing vLLM model: {model_id}")
111
- llm = LLM(model=model_id)
112
- tokenizer = AutoTokenizer.from_pretrained(model_id)
 
 
 
 
 
 
 
 
113
  sampling_params = SamplingParams(
114
  temperature=temperature,
115
  max_tokens=max_tokens,
 
25
 
26
  import polars as pl
27
  from datasets import Dataset, load_dataset
28
+ from huggingface_hub import login, dataset_info, snapshot_download
29
  from tqdm.auto import tqdm
30
  from transformers import AutoTokenizer
31
  from vllm import LLM, SamplingParams
 
106
  input_dataset_id, card_type, min_likes, min_downloads
107
  )
108
 
109
+ # Download model to local directory first
110
+ logger.info(f"Downloading model {model_id} to local directory...")
111
+ local_model_path = snapshot_download(
112
+ repo_id=model_id,
113
+ resume_download=True
114
+ )
115
+ logger.info(f"Model downloaded to: {local_model_path}")
116
+
117
+ # Initialize model and tokenizer from local path
118
+ logger.info(f"Initializing vLLM model from local path: {local_model_path}")
119
+ llm = LLM(model=local_model_path)
120
+ tokenizer = AutoTokenizer.from_pretrained(local_model_path)
121
  sampling_params = SamplingParams(
122
  temperature=temperature,
123
  max_tokens=max_tokens,