Spaces:
Sleeping
Sleeping
Commit
·
17da9d8
1
Parent(s):
2990eac
add
Browse files- app.py +5 -2
- requirements.txt +1 -0
app.py
CHANGED
@@ -7,12 +7,15 @@ def load_model(model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B"):
|
|
7 |
pipe = pipeline(
|
8 |
"text-generation",
|
9 |
model=model_name,
|
10 |
-
device_map="
|
11 |
torch_dtype=torch.float16,
|
12 |
trust_remote_code=True,
|
13 |
truncation=True,
|
14 |
max_new_tokens=2048,
|
15 |
-
model_kwargs={
|
|
|
|
|
|
|
16 |
)
|
17 |
return pipe
|
18 |
|
|
|
7 |
pipe = pipeline(
|
8 |
"text-generation",
|
9 |
model=model_name,
|
10 |
+
device_map="sequential",
|
11 |
torch_dtype=torch.float16,
|
12 |
trust_remote_code=True,
|
13 |
truncation=True,
|
14 |
max_new_tokens=2048,
|
15 |
+
model_kwargs={
|
16 |
+
"low_cpu_mem_usage": True,
|
17 |
+
"offload_folder": "offload"
|
18 |
+
}
|
19 |
)
|
20 |
return pipe
|
21 |
|
requirements.txt
CHANGED
@@ -2,4 +2,5 @@ streamlit
|
|
2 |
torch>=2.0.0
|
3 |
transformers>=4.35.0
|
4 |
accelerate>=0.20.0
|
|
|
5 |
gradio>=4.0.0
|
|
|
2 |
torch>=2.0.0
|
3 |
transformers>=4.35.0
|
4 |
accelerate>=0.20.0
|
5 |
+
safetensors>=0.3.1
|
6 |
gradio>=4.0.0
|