zhuohaoyu commited on
Commit
241089a
·
1 Parent(s): cd65843
Files changed (1) hide show
  1. app.py +4 -4
app.py CHANGED
@@ -33,23 +33,23 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
33
  from transformers.generation import GenerationConfig
34
 
35
 
36
- DEFAULT_CKPT_PATH = 'WisdomShell/CodeShell-7B-Chat'
37
 
38
 
39
  def _load_model_tokenizer(args):
40
  tokenizer = AutoTokenizer.from_pretrained(
41
- 'WisdomShell/CodeShell-7B-Chat', trust_remote_code=True, resume_download=True,
42
  )
43
 
44
  model = AutoModelForCausalLM.from_pretrained(
45
- 'WisdomShell/CodeShell-7B-Chat',
46
  trust_remote_code=True,
47
  resume_download=True,
48
  torch_dtype=torch.bfloat16
49
  ).eval()
50
 
51
  config = GenerationConfig.from_pretrained(
52
- 'WisdomShell/CodeShell-7B-Chat', trust_remote_code=True, resume_download=True,
53
  )
54
 
55
  return model, tokenizer, config
 
33
  from transformers.generation import GenerationConfig
34
 
35
 
36
+ DEFAULT_CKPT_PATH = 'WisdomShell/CodeShell-7B-Chat-int4'
37
 
38
 
39
  def _load_model_tokenizer(args):
40
  tokenizer = AutoTokenizer.from_pretrained(
41
+ 'WisdomShell/CodeShell-7B-Chat-int4', trust_remote_code=True, resume_download=True,
42
  )
43
 
44
  model = AutoModelForCausalLM.from_pretrained(
45
+ 'WisdomShell/CodeShell-7B-Chat-int4',
46
  trust_remote_code=True,
47
  resume_download=True,
48
  torch_dtype=torch.bfloat16
49
  ).eval()
50
 
51
  config = GenerationConfig.from_pretrained(
52
+ 'WisdomShell/CodeShell-7B-Chat-int4', trust_remote_code=True, resume_download=True,
53
  )
54
 
55
  return model, tokenizer, config