Spaces:
Running
on
CPU Upgrade
Running
on
CPU Upgrade
Commit
·
734af3d
1
Parent(s):
90d6ec0
load weights manually
Browse files- 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 |
-
#
|
110 |
-
logger.info(f"
|
111 |
-
|
112 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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,
|