Dart (Danbooru Tags Transformer) v1
This model is a fine-tuned Dart (Da nboor u T ags Transformer) model that generates danbooru tags.
If you are a developer and want to finetune, it's recommended using the base version,
p1atdev/dart-v1-base , instead
Usage
Using AutoModel
🤗 Transformers library is required.
pip install -U transformers1 import torch
2 from transformers import AutoTokenizer , AutoModelForCausalLM , GenerationConfig
3
4 MODEL_NAME = "p1atdev/dart-v1-sft"
5
6 tokenizer = AutoTokenizer . from_pretrained ( MODEL_NAME , trust_remote_code = True ) # trust_remote_code is required for tokenizer
7 model = AutoModelForCausalLM . from_pretrained ( MODEL_NAME , torch_dtype = torch . bfloat16 )
8
9 prompt = "<|bos|><rating>rating:sfw, rating:general</rating><copyright>original</copyright><character></character><general><|long|>1girl<|input_end|>"
10 inputs = tokenizer ( prompt , return_tensors = "pt" ) . input_ids
11
12 with torch . no_grad ( ) :
13 outputs = model . generate ( inputs , generation_config = model . generation_config )
14
15 print ( tokenizer . decode ( outputs [ 0 ] , skip_special_tokens = True ) )
16 # rating:sfw, rating:general, 1girl, ahoge, braid, closed eyes, collared dress, dress, flower, full body, hair flower, hair ornament, long hair, night, night sky, outdoors, parted lips, pink flower, pink hair, short sleeves, sky, solo, straight hair, sunflower, very long hair, white flower
You can use tokenizer.apply_chat_template to simplify constructiing of prompts:
1 inputs = tokenizer . apply_chat_template ( {
2 "rating" : "rating:sfw, rating:general" ,
3 "copyright" : "original" ,
4 "character" : "" ,
5 "general" : "1girl" ,
6 "length" : "<|long|>"
7 } , return_tensors = "pt" , tokenize = True ) # tokenize=False to preview prompt
8 # same as input_ids of "<|bos|><rating>rating:sfw, rating:general</rating><copyright>original</copyright><character></character><general><|long|>1girl<|input_end|>"
9 with torch . no_grad ( ) :
10 outputs = model . generate ( inputs , generation_config = generation_config )
See
chat_templating document for more detail about
apply_chat_template.
Flash attention (optional)
Using flash attention can optimize computations, but it is currently only compatible with Linux.
Accelerate with ORTModel
🤗 Optimum library is also compatible, for the high performance inference using ONNX.
pip install "optimum[onnxruntime]"Two ONNX models are provided:
Both can be utilized based on the following code:
1 import torch
2 from transformers import AutoTokenizer , GenerationConfig
3 from optimum . onnxruntime import ORTModelForCausalLM
4
5 MODEL_NAME = "p1atdev/dart-v1-sft"
6
7 tokenizer = AutoTokenizer . from_pretrained ( MODEL_NAME , trust_remote_code = True )
8
9 # normal version
10 ort_model = ORTModelForCausalLM . from_pretrained ( MODEL_NAME )
11
12 # qunatized version
13 # ort_model = ORTModelForCausalLM.from_pretrained(MODEL_NAME, file_name="model_quantized.onnx")
14
15 inputs = tokenizer . apply_chat_template ( {
16 "rating" : "rating:sfw, rating:general" ,
17 "copyright" : "original" ,
18 "character" : "" ,
19 "general" : "1girl" ,
20 "length" : "<|long|>"
21 } , return_tensors = "pt" , tokenize = True )
22
23 with torch . no_grad ( ) :
24 outputs = ort_model . generate ( inputs , generation_config = model . generation_config )
25
26 print ( tokenizer . decode ( outputs [ 0 ] , skip_special_tokens = True ) )
Prompt guide
Due to training with a specialized prompt format, natural language is not supported .
The trained sentences are essentially composed of the following elements, arranged in the strict order shown below:
<|bos|>: The bos (begin of sentence) token
<rating>[RATING_PARENT], [RATING_CHILD]</rating>: The block of rating tags
[RATING_PARENT]: rating:sfw, rating:nsfw
[RATING_CHILD]:
if [RATING_PARENT] is rating:sfw: rating:general, rating:sensitive
else: rating:questionable, rating:explicit
<copyright>[COPYRIGHT, ...]</copyright>: The block of copyright tags.
[COPYRIGHT, ...]: All supported copyright tags can be seen in here
<character>[CHARACTER, ...]</character>: The block of character tags.
[CHARACTER, ...]: All supported character tags can be seen in here
<general>[LENGTH_TOKEN][GENERAL, ...]<|input_end|>[COMPLETION]</general>: The block of general tags.
[LENGTH_TOKEN]: A token to specify total amount of general tags.
Avaiable:
<|very_short|>: less than 10 tags<|short|>: less than 20 tags<|long|>: less than 40 tags (recommended)<|very_long|>: more than 40 tags
[GENERAL, ...]: All supported general tags can be seen in here
<|input_end|>: A tag to show the end of input. Set this token at last of prompt.[COMPLETION]: The model complete tags in alphabetical order.
<|eos|>: The eos (end of sentence) token
Tags other than special tokens are separated by commas.
You can place tags in any order you like in each block.
Example sentence:
<|bos|><rating>rating:sfw, rating:general</rating><copyright>vocaloid</copyright><character>hatsune miku</character><general><|long|>solo, 1girl, very long hair<|input_end|>blue hair, cowboy shot, ...</general><|eos|>Therefore, to complete the tags, the input prompt should be as follows:
without any copyright and character tags
<|bos|><rating>rating:sfw, rating:general</rating><copyright></copyright><character></character><general><|very_long|>1girl, solo, cat ears<|input_end|>
specifing copyright and character tags
<|bos|><rating>rating:sfw, rating:general</rating><copyright>sousou no frieren</copyright><character>frieren</character><general><|long|>1girl, solo, from side<|input_end|>Model Details
Model Description
Developed by: Plat
Model type: Causal language model
Language(s) (NLP): Danbooru tags
License: Apache-2.0
Bias, Risks, and Limitations
Since this model is a pre-trained model, it cannot accommodate flexible specifications.
Training Details
Training Data
This model was trained with:
isek-ai/danbooru-tags-2023 : 6M size of danbooru tags dataset since 2005 to 2023
Only data from 2020 onwards was used for SFT.
Training Procedure
Trained using 🤗 transformers' trainer.
Preprocessing
Preprocessing was conducted through the following process:
Remove data where general tags is null.
Remove general tags that appear less than 100 times.
Remove undesirable tags such as watermark and bad anatomy.
Remove based on the number of tags attached to a single post (following rules):
Remove if more than 100 for general tags.
Remove if more than 5 for copyright tags.
Remove if more than 10 for character tags.
Remove posts created before 2020
Set length token according to each tags length
Shuffle some tags in the following rule:
Include people tags (e.g. 1girl, no humans) tags in the shuffle-group with a 95% probability, and do not do so with a 5% probability.
Get tags at a random percentage between 0% and 75% to create a shuffle-group.
Shuffle tags in shuffle-group and concatnate with <|input_end|> token and remains in alphabetical order.
Concatnate all categories
Training Hyperparameters
The following hyperparameters were used during training:
learning_rate: 0.0001
train_batch_size: 32
eval_batch_size: 8
seed: 42
gradient_accumulation_steps: 2
total_train_batch_size: 64
optimizer: Adam with betas=(0.9,0.999) and epsilon=1e-08
lr_scheduler_type: linear
lr_scheduler_warmup_steps: 500
num_epochs: 1
Evaluation
Evaluation has not been done yet and it needs to evaluate.
Technical Specifications
Model Architecture and Objective
The architecture of this model is
OPT (Open Pretrained Transformer) , but the position embeddings was not trained.
Compute Infrastructure
In house
Hardware
1x RTX 3070 Ti
Software
Dataset processing: 🤗 Datasets
Training: 🤗 Transformers
Optimizing: 🤗 Optimum
SFT: 🤗 TRL
More Information [optional]
[More Information Needed]
