Update README.md
Browse files
README.md
CHANGED
@@ -22,10 +22,10 @@ import torch
|
|
22 |
|
23 |
REPO_NAME = 'schuler/experimental-JP47D20'
|
24 |
|
25 |
-
def load_model(
|
26 |
-
tokenizer = LlamaTokenizer.from_pretrained(
|
27 |
-
generator_conf = GenerationConfig.from_pretrained(
|
28 |
-
model = AutoModelForCausalLM.from_pretrained(
|
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 |
|