schuler commited on
Commit
931f5e1
·
verified ·
1 Parent(s): b1aa907

Update README.md

Browse files
Files changed (1) hide show
  1. README.md +4 -4
README.md CHANGED
@@ -22,10 +22,10 @@ import torch
22
 
23
  REPO_NAME = 'schuler/experimental-JP47D20'
24
 
25
- def load_model(repo_name):
26
- tokenizer = LlamaTokenizer.from_pretrained(REPO_NAME, trust_remote_code=True)
27
- generator_conf = GenerationConfig.from_pretrained(REPO_NAME)
28
- model = AutoModelForCausalLM.from_pretrained(REPO_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="eager")
29
  # model.to('cuda')
30
  return tokenizer, generator_conf, model
31
 
 
22
 
23
  REPO_NAME = 'schuler/experimental-JP47D20'
24
 
25
+ def load_model(local_repo_name):
26
+ tokenizer = LlamaTokenizer.from_pretrained(local_repo_name, trust_remote_code=True)
27
+ generator_conf = GenerationConfig.from_pretrained(local_repo_name)
28
+ model = AutoModelForCausalLM.from_pretrained(local_repo_name, trust_remote_code=True, torch_dtype=torch.bfloat16, attn_implementation="eager")
29
  # model.to('cuda')
30
  return tokenizer, generator_conf, model
31