Tonic commited on
Commit
ca44a43
1 Parent(s): f58e1e9

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +6 -1
app.py CHANGED
@@ -12,7 +12,12 @@ description = "You can use [🐋🐳microsoft/Orca-2-13b](https://huggingface.co
12
  # os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:24'
13
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
14
  model_name = "microsoft/Orca-2-13b"
15
- model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto')
 
 
 
 
 
16
  tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False,)
17
 
18
  class OrcaChatBot:
 
12
  # os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:24'
13
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
14
  model_name = "microsoft/Orca-2-13b"
15
+ offload_folder = './model_weights'
16
+
17
+ if not os.path.exists(offload_folder):
18
+ os.makedirs(offload_folder)
19
+
20
+ model = AutoModelForCausalLM.from_pretrained(model_name, device_map='auto', offload_folder=offload_folder)
21
  tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=False,)
22
 
23
  class OrcaChatBot: