kishkath commited on
Commit
767ab2c
·
verified ·
1 Parent(s): 8da2345

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +5 -13
app.py CHANGED
@@ -19,21 +19,13 @@ tokenizer = AutoTokenizer.from_pretrained(model_dir)
19
  # offload_folder=offload_dir # Directory to offload layers (for larger models).
20
  # )
21
 
22
- ## CPU
23
- # base_model = AutoModelForCausalLM.from_pretrained(
24
- # model_dir,
25
- # device_map="cpu", # Force CPU usage
26
- # torch_dtype=torch.float32, # Use float32 for CPU
27
- # trust_remote_code=True,
28
- # offload_folder=offload_dir # Directory to offload layers (for larger models).
29
- # )
30
-
31
- ## GPU
32
  base_model = AutoModelForCausalLM.from_pretrained(
33
  model_dir,
34
- device_map="auto", # Automatically uses available GPU(s)
35
- torch_dtype=torch.float16, # Use half precision for faster inference on GPU
36
- trust_remote_code=True
 
37
  )
38
 
39
 
 
19
  # offload_folder=offload_dir # Directory to offload layers (for larger models).
20
  # )
21
 
22
+ # CPU
 
 
 
 
 
 
 
 
 
23
  base_model = AutoModelForCausalLM.from_pretrained(
24
  model_dir,
25
+ device_map="cpu", # Force CPU usage
26
+ torch_dtype=torch.float32, # Use float32 for CPU
27
+ trust_remote_code=True,
28
+ offload_folder=offload_dir # Directory to offload layers (for larger models).
29
  )
30
 
31