Spaces:
Build error
Build error
File size: 1,343 Bytes
e7d3e35 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
#!/usr/bin/env python
# This script creates a super tiny model that is useful inside tests, when we just want to test that
# the machinery works, without needing to check the quality of the outcomes.
#
# usage: adjust the configs if wanted, but otherwise just run the script
from pathlib import Path
from transformers import LlamaConfig, LlamaForCausalLM, LlamaTokenizer
mname_tiny = "tiny-random-LlamaForCausalLM"
path = Path(mname_tiny)
path.mkdir(parents=True, exist_ok=True)
config = LlamaConfig()
config.update(
dict(
vocab_size=32000,
hidden_size=16,
intermediate_size=16 * 4,
num_hidden_layers=2,
num_attention_heads=4,
)
)
model = LlamaForCausalLM(config)
tokenizer = LlamaTokenizer.from_pretrained("path_to_llama_7b")
# Test w/ one text
query = "This is a test"
query_tokens = tokenizer(query, return_tensors="pt")
input = {
"input_ids": query_tokens["input_ids"],
"attention_mask": query_tokens["attention_mask"],
}
out_gen = model.generate(**input)
text = tokenizer.batch_decode(out_gen)
# Save model + config + tokenizer
model.half() # makes it smaller
model.save_pretrained(path)
tokenizer.save_pretrained(path)
# test we can load it back
model = LlamaForCausalLM.from_pretrained(path)
print(f"Generated {mname_tiny} - Upload the generated folder to the hub")
|