はじめに
TinySwallow-1.5Bは、新手法「TAID」を用いた小規模日本語言語モデルです。本記事では、Windows 11 + WSL2環境での構築手順と、Web UIの実装方法を解説します。
環境要件
- Windows 11（WSL2対応）
- Ubuntu 24.04 on WSL2
- NVIDIA GPU（RTX 40シリーズ推奨）
- NVIDIA Driver 535以上
1. WSL2環境の準備
# PowerShellを管理者権限で実行
wsl --install -d Ubuntu-24.04
wsl --set-version Ubuntu-24.04 2
2. 基本環境のセットアップ
# パッケージの更新
sudo apt update && sudo apt upgrade -y
# 必要なパッケージのインストール
sudo apt install -y python3.11 python3.11-venv python3.11-dev build-essential ninja-build cmake
# 仮想環境の作成
mkdir ~/projects && cd ~/projects
python3.11 -m venv tinyswallow-env
source tinyswallow-env/bin/activate
3. PyTorch環境の構築
# 基本パッケージのインストール
pip install --upgrade pip wheel setuptools
# PyTorchのインストール（CUDA 12.1対応）
pip install torch==2.5.0 torchvision==0.20.0 --index-url https://download.pytorch.org/whl/cu121
# 動作確認
python -c "import torch; print(f'CUDA available: {torch.cuda.is_available()}')"
4. Flash Attentionのセットアップ
# ソースからビルド
cd ~
git clone https://github.com/Dao-AILab/flash-attention
cd flash-attention
git checkout v2.7.3
# ビルド設定
export TORCH_CUDA_ARCH_LIST="8.9" # RTX 40シリーズ用
MAX_JOBS=4 pip install . --no-build-isolation --verbose
5. TinySwallow-1.5B環境の構築
pip install transformers==4.40.0 accelerate==0.29.3 sentencepiece==0.2.0
# 動作確認スクリプト
python -c """
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
model = AutoModelForCausalLM.from_pretrained(
'SakanaAI/TinySwallow-1.5B-Instruct',
device_map='auto',
torch_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained('SakanaAI/TinySwallow-1.5B-Instruct')
prompt = 'AIについて説明してください。'
inputs = tokenizer(prompt, return_tensors='pt').to(model.device)
outputs = model.generate(**inputs, max_new_tokens=200)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))
"""
6. Web UIの実装
# Gradioのインストール
pip install gradio
# app.pyの作成
cat << 'EOF' > app.py
import gradio as gr
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
def generate_text(prompt, max_length=200, temperature=0.7):
inputs = tokenizer(prompt, return_tensors="pt").to(device)
outputs = model.generate(
**inputs,
max_new_tokens=max_length,
temperature=temperature,
do_sample=True
)
return tokenizer.decode(outputs[0], skip_special_tokens=True)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
"SakanaAI/TinySwallow-1.5B-Instruct",
device_map="auto",
torch_dtype=torch.bfloat16
)
tokenizer = AutoTokenizer.from_pretrained("SakanaAI/TinySwallow-1.5B-Instruct")
iface = gr.Interface(
fn=generate_text,
inputs=[
gr.Textbox(label="プロンプト"),
gr.Slider(minimum=50, maximum=500, value=200, label="最大生成トークン数"),
gr.Slider(minimum=0.1, maximum=1.0, value=0.7, label="温度")
],
outputs=gr.Textbox(label="生成テキスト"),
title="TinySwallow-1.5B Chat",
description="日本語に特化した小規模言語モデル"
)
if __name__ == "__main__":
iface.launch(server_name="0.0.0.0", server_port=7861)
EOF
# Web UIの起動
python app.py
7. 動作確認
- ブラウザで
http://localhost:7861にアクセス - プロンプトを入力して生成を実行
- 生成されたテキストを確認
トラブルシューティング
CUDA関連
-
nvidia-smiでGPUが認識されない場合# WSL2側のドライバ確認 ls /usr/lib/wsl/lib/nvidia-smi
ポート関連
- ポート7861が使用中の場合
# 使用中のポートを確認 lsof -i :7861 # 別のポートを使用 export GRADIO_SERVER_PORT=7862
メモリ関連
- VRAM不足の場合
# 4bit量子化を有効化 model = AutoModelForCausalLM.from_pretrained( "SakanaAI/TinySwallow-1.5B-Instruct", device_map="auto", load_in_4bit=True )
参考文献
ライセンス
本記事で使用しているコードは、MITライセンスで提供されています。
著者
タグ: Python, AI, WSL2, CUDA, PyTorch, Transformers, Web UI, TinySwallow