blockenters commited on
Commit
17da9d8
·
1 Parent(s): 2990eac
Files changed (2) hide show
  1. app.py +5 -2
  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="balanced_low_0",
11
  torch_dtype=torch.float16,
12
  trust_remote_code=True,
13
  truncation=True,
14
  max_new_tokens=2048,
15
- model_kwargs={"low_cpu_mem_usage": True}
 
 
 
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