初めに
With the recent release of #TinyLlama, SLMs have attracted a lot of attention. I re-released my previously trained SLM - LiteLlama under the MIT license, which has 460M parameters trained with 1T tokens. I hope to contribute a bit to the community.https://t.co/NRQPv6Jlr9
— Xiaotian (Max) Han (@XiaotianHan1) 2024年1月7日
400MB程度のSLMが公開されたので、触っていきます
環境
- L4 GPU
- ubuntu22.04
- Jupyter Notebook
準備
以下のライブラリを入れます
pip install torch pip install transformers
実行
モデルのロード
import torch from transformers import AutoTokenizer, AutoModelForCausalLM model_path = 'ahxt/LiteLlama-460M-1T' model = AutoModelForCausalLM.from_pretrained(model_path) tokenizer = AutoTokenizer.from_pretrained(model_path) model.eval()
サンプルプロンプト
推論
prompt = 'Q: What is the largest bird?\nA:' input_ids = tokenizer(prompt, return_tensors="pt").input_ids tokens = model.generate(input_ids, max_length=20) print( tokenizer.decode(tokens[0].tolist(), skip_special_tokens=True) )
結果
Q: What is the largest bird? A: The largest bird is the largest bird.
まどマギ
推論
prompt = 'まどマギで一番可愛いキャラは？' input_ids = tokenizer(prompt, return_tensors="pt").input_ids tokens = model.generate(input_ids, max_length=100) print( tokenizer.decode(tokens[0].tolist(), skip_special_tokens=True) )
結果
まどマギで一番可愛いキャラは？２．２．２．２．２．２．２．２．２．２．２．２．２．�
まどマギQA
推論
prompt = 'Q.まどマギで一番可愛いキャラは?\nA:' input_ids = tokenizer(prompt, return_tensors="pt").input_ids tokens = model.generate(input_ids, max_length=100) print( tokenizer.decode(tokens[0].tolist(), skip_special_tokens=True) )
結果
Q.まどマギで一番可愛いキャラは? A: これはそれはそれは? A: これはそれはそれは? A: これはそれはそれは? A: これはそれはそれは? A: これはそれはそれは
