transformers 4.39.0がmambaをサポートしたので、kotomamba-2.8B-v1.0をtransformersで動かしてみた。
#! /usr/bin/python3
from transformers import LlamaTokenizerFast,MambaConfig,MambaForCausalLM,TextGenerationPipeline
from transformers.utils import cached_file
tkz=LlamaTokenizerFast(vocab_file=cached_file("kotoba-tech/kotomamba-2.8B-v1.0","tokenizer.model"),add_bos_token=False,pad_token="<pad>")
cfg=MambaConfig.from_pretrained("kotoba-tech/kotomamba-2.8B-v1.0",hidden_size=2560,intermediate_size=5120,num_hidden_layers=64,time_step_rank=160,vocab_size=96872,bos_token_id=1,eos_token_id=2,pad_token_id=4)
mdl=MambaForCausalLM.from_pretrained("kotoba-tech/kotomamba-2.8B-v1.0",config=cfg)
mdl.resize_token_embeddings(len(tkz))
tgn=TextGenerationPipeline(model=mdl,tokenizer=tkz,max_new_tokens=128)
nlp=lambda txt:tgn(txt)[0]["generated_text"]
print(nlp("国境の長いトンネルを抜けると雪国であった。夜の底が白くなった。"))
MambaConfigの値に無理矢理感がある上に、LlamaTokenizerFastを借りていたりするし、その結果、トークナイザの単語数とモデルのtoken_embeddingsが合わなくなってしまって、resize_token_embeddingsとか恐ろしいことをしている。何とか『雪国』冒頭部の続きを生成させてみたところ、私(安岡孝一)の手元では以下の結果が出力された。
国境の長いトンネルを抜けると雪国であった。夜の底が白くなった。信号所に汽車が止まった。
- 1945年8月15日、終戦の日。
- 1945年8月15日、終戦の日。
- 1945年8月15日、終戦の日。
- 1945年8月15日、終戦の日。
- 1945年8月15日、終戦の日。
- 1945年8月15日、終戦の日。
- 19
「信号所に汽車が止まった。」が生成されているので、kotomamba-2.8B-v1.0は、この一節を知っているらしい。どうしてそれが「終戦の日」に繫がるのかは謎だが、そういうものなのだろう。さて、どういう形で継続学習させたらいいかな。