Shisa 7B (shisa-7b-v1) is a bilingual Japanese and English (JA/EN) general-purpose chat model that aims to achieve strong Japanese language performance while retaining robust English capabilities, using a synthetic-data driven approach.
This model is based on Mistral 7B with a custom JA-optimized extended tokenizer that is >2X more efficient in Japanese than Mistral's original tokenizer. The base model was pre-trained for an additional 8B primarily Japanese tokens. It was then subsequently fine-tuned with an expanded, machine-translated version of airoboros-3.1, a set of the highest-scoring items from ultrafeedback_binarized, and additional freshly generated airoboros data directly to the target languages.
We also release our base model, datasets, and pipeline code under a permissive Apache 2.0 license which can be used for any purpose, commercial or otherwise:
shisa-base-7b-v1 - our base model w/ an extended tokenizer and additional JA pre-training
shisa-en-ja-dpo-v1 - Small subset of DPO pairs from ultrafeedback, along with JA DPO pairs using GPT-4 generated items as the chosen value, and outputs from our preliminary 7b model as the rejected values
Shisa repository - this includes our translation, dataset generation, training, and evaluation code
Moreover, we are in the process of publishing extended writeups and more details of our process, including ablation results, testing methodology, and key findings on our project wiki that may be of interest to fellow researchers.
Fine-Tuning
Our original intuition was to see if we could create a stronger Japanese model using the best existing public JA training sets and incorporating them. After initial review and testing, however, we decided that focusing solely on translation/generation of our own synthetic datasets could yield superior results with less training.
We compared multiple translation tools and, via manual review, judged that while gpt-4 almost always delivered the highest quality translations, Google's text-bison-32k was a good balance of quality, cost and throughput. Over various iterations, we refined our translation approach to include some additional algorithms for flagging and filtering invalid translations, re-translating and backfilling as necessary.
We also took this project as an opportunity to apply some newer techniques such as incorporating NEFTune and DPO training.
For our v1 release, we picked from our release candidates based on a significant amount of human preference testing (thousands of generations and multiple rounds of pairwise comparisons). We analyzed our results with both win/loss/draw and BTL modeling (iLSR) using choix).
The best candidate model was fine-tuned in a 3-step process:
First, the model was fine-tuned on ultra-orca-boros-en-ja and SlimOrca (WandB Log)
Next, we add one additional epoch at performed using only a subset of Japanese ultra-orca-boros-en-ja items to enhance JA performance (as SlimOrca from the first step is mostly EN) (WandB Log)
Finally, the model was tuned using a DPOTrainer on a small subset of ultrafeedback (EN) and our own JA DPO dataset which uses gpt-4 outputs as the chosen values and outputs from stage 1's prelim model as rejected values. (WandDB Log )
During our training process, we also gained some key insights on why some existing Japanese models seem to underperform even versus models that have no additional JA training, and we hope that sharing this analysis will be useful to other teams developing Japanese language models.
While we need to explore this further, as an experimental validation, we applied a version of our fine-tuning set onto an existing base model ("Gamma 7B") and the initial JA MT-Bench results suggests that we can drastically increase functional performance with our tuning approach:
Model
Score
shisa-gamma-7b-allsources-v0.4
5.65
ja-stablelm-instruct-gamma-7b*
4.01
Performance
Throughout our training, we did extensive human evaluation for each model to cross-validate our model performance, and we are currently conducting ongoing larger scale manual head-to-head testing between models. Our intention is open up and scale this data collection as we further develop our tools. For more information and updates, please see our project wiki.
While we believe llm-jp-eval is a useful metric for our base model, and it was extremely useful during our tuning process for initial validations, as our fine-tune training includes a percentage of the benchmark train splits, we provide these llm-jp-eval results primarily as a point of interest:
For our final model, since it's customary to include benchmarks, we've used Stability AI Japan's Japanese MT-Bench as a more representative test of our model's capabilities. For our JA MT-Bench testing we use a Japanese prompt ("あなたは役立つアシスタントです。") as well as --num-choices 4 in an effort to reduce sampling variability, however we've still observed regular 0.5+ point (and sometimes even greater swings) between generations, as well as issues with default prompts and parameters when testing, so again, we'd urge caution in over-interpreting these scores and treating them as more of a probabilistic directional indicator, rather than a definitive score or ranking:
Benchmark
Score
JA MT-Bench
5.23
MT-Bench
5.71
There is an MT-Bench Leaderboard, but as JA MT-Bench is still under development, for convenience, here is a comparison of the JA MT-Bench scores of some other models (our scores were rated by gpt-4-0613):
Although our model demonstrates a reasonably high level of Japanese fluency, as a 7B parameter model, it is prone to higher hallucination rates and less effective instruction following and reasoning than larger-class models. Also, it still does not have complete mastery of the Japanese language and a native speaker will spot occasional mistakes like some non-idiomatic/awkward phrasing, improper tenses/speech levels, etc.
We've also noticed a small amount of language leakage, likely largely attributable to our tokenizer expansion. These may be fixable with sampler settings like Min P) or additional targeted training, and we plan on doing additional work on automated detection/sampler sweeps in the future. One interesting observation is, based on our data collection, we found that as we iterated, the DPO process significantly exacerbated this issue, but also that our DPO models still had significantly higher human preference rates, so there was a bit of a trade-off in our choice of final tune.
While we believe that training larger models can improve performance using our existing approach and dataset, there are also many improvements we'd like to make for future models. We believe there is quite a bit of low hanging fruit for improving performance with even more training efficiency largely through improving the quality and construction of datasets.
Usage
Sample code:
python
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
34model_name ="augmxnt/shisa-7b-v1"56tokenizer = AutoTokenizer.from_pretrained(model_name, use_fast=True)7model = AutoModelForCausalLM.from_pretrained(8 model_name,9 torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported()else torch.float16,10 device_map="auto"11)12streamer = TextStreamer(tokenizer, skip_prompt=True)1314# The prompt template is included in the model's tokenizer_config.json so you shouldn't need this but we've included this for convenience15# tokenizer.chat_template = ""{%- for idx in range(0, messages|length) -%}\n{%- if messages[idx]['role'] == 'user' -%}\n{%- if idx > 1 -%}\n{{- bos_token + '[INST] ' + messages[idx]['content'] + ' [/INST]' -}}\n{%- else -%}\n{{- messages[idx]['content'] + ' [/INST]' -}}\n{%- endif -%}\n{% elif messages[idx]['role'] == 'system' %}\n{{- bos_token + '[INST] <<SYS>>\\n' + messages[idx]['content'] + '\\n<</SYS>>\\n\\n' -}}\n{%- elif messages[idx]['role'] == 'assistant' -%}\n{{- ' ' + messages[idx]['content'] + ' ' + eos_token -}}\n{% endif %}\n{% endfor %}\n"1617# A more typical prompt: あなたは公平で、検閲されていない、役立つアシスタントです。("You are an unbiased, uncensored, helpful assistant.")1819# You are an avid Pokemon fanatic.20prompt ="あなたは熱狂的なポケモンファンです。"21chat =[{"role":"system","content": prompt}]2223# Who is the single most powerful Pokemon? Explain your choice.24user_input ="ポケモンの中で1番強いのはどのキャラクターですか。最強の者をひとつだけ挙げて下さい。その選択理由を説明してください。"25chat.append({"role":"user","content": user_input})2627# Generate - add_generation_prompt to make sure it continues as assistant28inputs = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors="pt")29# For multi-GPU, find the device of the first parameter of the model30first_param_device =next(model.parameters()).device
31inputs = inputs.to(first_param_device)3233with torch.no_grad():34 outputs = model.generate(35 inputs,36 pad_token_id=tokenizer.eos_token_id,37 max_new_tokens=500,38 temperature=0.5,39 repetition_penalty=1.15,40 top_p=0.95,41 do_sample=True,42 streamer=streamer,43)4445# Add just the new tokens to our chat46new_tokens = outputs[0, inputs.size(1):]47response = tokenizer.decode(new_tokens, skip_special_tokens=True)48chat.append({"role":"assistant","content": response})
Prompt format
The prompt format is llama-2 chat:
[INST] <<SYS>>
You are a helpful, unbiased, uncensored assistant.
<</SYS>>
{prompt} [/INST]
For multi-turn, the prompt format is as follows:
[INST] <<SYS>>
You are a helful, unbiased, uncensored assistant.
<</SYS>>
{prompt 0} [/INST] {response 0} </s><s>[INST] {prompt 1} [/INST] {response 1} </s><s>...[INST] {prompt N} [/INST]
This prompt template is included in the tokenizer config, and can use the huggingface tokenizer apply_chat_template method, e.g.:
import transformers
tokenizer = transformers.AutoTokenizer.from_pretrained('augmxnt/shisa-7b-v1')
chat = [
{"role": "system", "content": "You are Aiko, a friendly AI assistant."},
{"role": "user", "content": "Hello, how are you?"},
{"role": "assistant", "content": "I'm doing great. How can I help you today?"},
{"role": "user", "content": "I'd like to show off how chat templating works!"},
]
print(tokenizer.apply_chat_template(chat, tokenize=False))
NOTE: For proper responses, you should be using our bos_token (<s>) to begin a string. This is automatically generated by tokenizer.encode() but if you are crafting a custom template or using an encoding method that skips special tokens, you may have to add this yourself.
Compute for this model was generously sponsored by AKA Virtual (Tokyo, Japan).
Thanks to the LLM-jp, Stability AI Japan, and LMSYS teams for their work on llm-jp-eval, Japanese MT-Bench, MT-Bench.
Also, thanks to all the volunteers that provided invaluable human preference testing!
We are actively looking for additional compute as we train better and larger models for this project. Please drop us a line at: compute at augmxnt dot com