crystal-technologies commited on
Commit
087de2e
·
1 Parent(s): 21c32dd

Update Perceptrix/engine.py

Browse files
Files changed (1) hide show
  1. Perceptrix/engine.py +4 -1
Perceptrix/engine.py CHANGED
@@ -4,8 +4,11 @@ from utils import setup_device
4
  import transformers
5
  import torch
6
  import tqdm
 
7
 
8
- model_name = "mosaicml/mpt-30b-instruct"
 
 
9
 
10
  config = transformers.AutoConfig.from_pretrained(
11
  model_name, trust_remote_code=True)
 
4
  import transformers
5
  import torch
6
  import tqdm
7
+ import os
8
 
9
+ model_name = os.environ.get('LLM_MODEL')
10
+
11
+ model_path = "models/CRYSTAL-instruct" if model_name != None else model_name
12
 
13
  config = transformers.AutoConfig.from_pretrained(
14
  model_name, trust_remote_code=True)