1. 導入
Pythonを愛用しているあなた、「この子、ちょっと遅いかも…」と感じたことがあるでしょう。
モデルの学習にやたら時間がかかったり、高パフォーマンスが求められる場面で「C++か…はぁ…」とため息をつく経験はありませんか？
そこで登場するのがMojo🔥
「Pythonの親戚みたいな顔してるくせにC++級のスピード!？ Rust級のメモリ安全性!?なんだコイツは…」と、すでに業界をザワつかせている新顔です。
Mojoは、Pythonのようなシンプルで読みやすい文法を保ちつつ、AI/MLやハイパフォーマンス計算をバリバリこなすことを目指して設計された言語です。
さて、このMojoが見せる魔法の正体に迫ってみましょう👀
2. Mojoの特徴
-
Pythonとの親和性
Pythonに似た構文とエコシステム統合が可能で、NumPyやPyTorchを直接活用でき、移行もスムーズ -
静的型付けとJITコンパイル
fn関数で型を明示すれば、コンパイル時チェックとJITコンパイルの相乗効果で高速実行が可能が可能 -
所有権モデルによるメモリ安全性
Rust風の所有権モデルにより、メモリ管理の問題をコンパイル時に解決。より安全かつ堅牢なコードを実現 -
AI/ML特化設計
高速な行列演算、GPU/TPUとの連携など、AI/ML分野で威力を発揮。BERTモデルの実装もスムーズ
3. Mojoの基本文法
MojoはPythonライクな書き心地に加え、静的型付けやメモリ安全性を組み合わせることで、高速かつ堅牢なコードが書けます。
以下では簡易的な文法例や特徴を紹介します。
3.0. Python統合
from python import PythonでPythonモジュールをインポートし、Pythonのライブラリをそのまま利用可能。
NumPyやPyTorchといったAIライブラリをMojoコードに直接組み込むことができます。これ最強!!
from python import Python
def main():
var np = Python.import_module("numpy")
var arr = np.array([1, 2, 3])
print(arr)
3.1. 関数定義と型の宣言
関数定義
Mojoには2種類の関数定義があります。
-
def: Pythonスタイル- 型注釈、宣言が不要で、素早く関数を書ける
- 実行時エラーが起きうるが、プロトタイプや小規模スクリプトに便利
-
fn: 静的型付けスタイル- 引数や戻り値に型を明示することで、コンパイル時チェックが働き、型安全性・高速性が向上
- RustやC++に近い感覚で安全なコードを記述できる
# def関数（型指定不要、動的）
def greet(name):
return "Hello, " + name + "!"
# fn関数（静的型付け必須）
fn greet2(name: String) -> String:
return "Hello, " + name + "!"
型の宣言
fn 関数を使用することで、以下の型を明示的に指定できます。
- 基本型: Int, Float, String, Bool など
- コンテナ型: List[T], Dict[K, V]
fn calculate_area(radius: Float) -> Float:
var pi: Float = 3.14159
return pi * radius * radius
var age: Int = 25
var name: String = "Alice"
var is_valid: Bool = True
3.2. エントリーポイント：main()関数
.mojoファイルをスクリプトとして実行する場合、main()関数が必要です。
このmain()がプログラムの開始地点となります。
def main():
print("Mojo is starting up!")
※Jupyter上ではmain() は不要で、対話的にMojoコードを実行可能。
3.3. 制御構文・ブロック
Python同様、インデントでブロックを表します。
if、for、whileなどはほぼPythonと同じ使い心地。
def loop_example():
for i in range(5):
if i % 2 == 0:
print(i, "is even")
else:
print(i, "is odd")
3.4. 構造体（struct）
structを用いることで、静的かつ高速なデータ型を定義可能。
Pythonのclassに似ていますが、Mojoのstructはコンパイル時に確定するため、動的な変更は不可。その代わり高速で安全。
struct MyPair:
var first: Int
var second: Int
fn __init__(out self, first: Int, second: Int):
self.first = first
self.second = second
fn dump(self):
print(self.first, self.second)
fn use_pair():
var p = MyPair(10, 20)
p.dump() # "10 20"と表示
3.5. トレイト（trait）
トレイトは、型に「これらのメソッドを必ず持っていてください」と約束させる仕組みです。
トレイトを適用すると、その型は決められたメソッドを必ず実装する必要があります。これにより、異なる型を共通の方法で扱えるようになります。
# トレイト定義
trait Printable:
fn print_self(self)
# トレイトに準拠する型
struct MyNum(Printable):
var value: Int
fn __init__(out self, val: Int):
self.value = val
fn print_self(self):
print("Value:", self.value)
# トレイトに基づくジェネリック関数
fn print_anything[T: Printable](item: T):
item.print_self()
# 使用例
fn main():
var number = MyNum(42)
print_anything(number) # 出力: Value: 42
4. BERT実装例
ここでは、PythonとMojoでBERTモデルを使った推論を比較しました。シンプルな推論タスクで、入力テキスト中の[MASK]トークンを予測する例です。実際のコードについては、詳細を以下のGitHubリポジトリで確認してください。
Pythonコード
import os
import platform
import signal
from argparse import ArgumentParser
import torch
from max import engine
from max.dtype import DType
from transformers import BertTokenizer
os.environ["TRANSFORMERS_VERBOSITY"] = "critical"
os.environ["TOKENIZERS_PARALLELISM"] = "false"
BATCH = 1
SEQLEN = 128
DEFAULT_MODEL_PATH = "../../models/bert-mlm.torchscript"
DESCRIPTION = "BERT model"
HF_MODEL_NAME = "bert-base-uncased"
def execute(model_path, text, input_dict):
"""
モデルをロードし、入力テキストを処理して、推論を実行し、マスクされたトークンを予測。
引数:
model_path (str): 学習済みモデルファイルのパス。
text (str): `[MASK]`トークンを含む入力テキスト。
input_dict (dict): 入力テンソルの仕様を定義する辞書。
戻り値:
str: 入力テキスト中の`[MASK]`に対応する予測トークン。
"""
session = engine.InferenceSession()
input_spec_list = [
engine.TorchInputSpec(shape=tensor.size(), dtype=DType.int64)
for tensor in input_dict.values()
]
model = session.load(model_path, input_specs=input_spec_list)
tokenizer = BertTokenizer.from_pretrained(HF_MODEL_NAME)
print("Processing input...")
inputs = tokenizer(
text,
return_tensors="pt",
padding="max_length",
truncation=True,
max_length=SEQLEN,
)
print("Input processed.\n")
masked_index = (inputs["input_ids"] == tokenizer.mask_token_id).nonzero(
as_tuple=True
)[1]
outputs = model.execute_legacy(**inputs)["result0"]
logits = torch.from_numpy(outputs[0, masked_index, :])
predicted_token_id = logits.argmax(dim=-1)
predicted_tokens = tokenizer.decode(
[predicted_token_id],
skip_special_tokens=True,
clean_up_tokenization_spaces=True,
)
return predicted_tokens
def main():
"""
コマンドライン引数を解析し、モデルを用いてマスクされたトークンの予測。
コマンドライン引数:
--text: `[MASK]`トークンを含む入力テキスト。
--model-path: 学習済みモデルファイルのパス。
出力:
入力テキストと`[MASK]`に対する予測結果。
"""
parser = ArgumentParser(description=DESCRIPTION)
parser.add_argument(
"--text",
type=str,
metavar="<text>",
required=True,
help="Masked language model.",
)
parser.add_argument(
"--model-path",
type=str,
default=DEFAULT_MODEL_PATH,
help="Directory for the downloaded model.",
)
args = parser.parse_args()
if "Intel" in platform.processor():
os.environ["OMP_NUM_THREADS"] = "1"
os.environ["MKL_NUM_THREADS"] = "1"
signal.signal(signal.SIGINT, signal.SIG_DFL)
torch.set_default_device("cpu")
input_dict = {
"input_ids": torch.zeros((BATCH, SEQLEN), dtype=torch.int64),
"attention_mask": torch.zeros((BATCH, SEQLEN), dtype=torch.int64),
"token_type_ids": torch.zeros((BATCH, SEQLEN), dtype=torch.int64),
}
outputs = execute(args.model_path, args.text, input_dict)
print(f"入力テキスト: {args.text}")
print(f"予測結果: {args.text.replace('[MASK]', outputs)}")
if __name__ == "__main__":
main()
Pythonの実行結果
input text: Tokyo is the [MASK] of Japan.
filled mask: Tokyo is the capital of Japan.
magic run bash run.sh 2.50s user 0.69s system 29% cpu 10.643 total
Mojoのコード
from max.engine import InputSpec, InferenceSession, Model
from pathlib import Path
from python import Python, PythonObject
from max.tensor import Tensor, TensorSpec
import sys
def execute(model: Model, text: String, transformers: PythonObject) -> String:
"""
モデルを使用してマスクされたトークンを予測します。
Args:
model (Model): ロード済みのモデル。
text (String): `[MASK]`を含む入力テキスト。
transformers (PythonObject): Hugging Faceのトークナイザーを含むライブラリ。
Returns:
String: 予測されたトークン。
"""
output_spec = model.get_model_output_metadata()[0]
max_seqlen = output_spec[1].value()
tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer(
text=text,
add_special_tokens=True,
padding="max_length",
truncation=True,
max_length=max_seqlen,
return_tensors="np",
)
input_ids = inputs["input_ids"]
token_type_ids = inputs["token_type_ids"]
attention_mask = inputs["attention_mask"]
outputs = model.execute(
"input_ids",
input_ids,
"token_type_ids",
token_type_ids,
"attention_mask",
attention_mask,
)
logits = outputs.get[DType.float32]("result0")
mask_idx = -1
for i in range(len(input_ids[0])):
if input_ids[0][i] == tokenizer.mask_token_id:
mask_idx = i
predicted_token_id = argmax(logits)[mask_idx]
return str(
tokenizer.decode(
predicted_token_id,
skip_special_tokens=True,
clean_up_tokenization_spaces=True,
)
)
def argmax(t: Tensor) -> List[Int]:
"""
テンソルの各列ごとに最大値のインデックスを計算します。
Args:
t (Tensor): 3次元のテンソル。
Returns:
List[Int]: 各列の最大値のインデックスを持つリスト。
"""
var res = List[Int](capacity=t.dim(1))
for i in range(t.dim(1)):
var max_val = Scalar[t.type].MIN
var max_idx = 0
for j in range(t.dim(2)):
if t[0, i, j] > max_val:
max_val = t[0, i, j]
max_idx = j
res.append(max_idx)
return res
def load_model(session: InferenceSession) -> Model:
"""
モデルを指定された入力仕様でロードします。
Args:
session (InferenceSession): 推論セッション。
Returns:
Model: ロードされたモデル。
"""
batch = 1
seqlen = 128
input_ids_spec = TensorSpec(DType.int64, batch, seqlen)
token_type_ids_spec = TensorSpec(DType.int64, batch, seqlen)
attention_mask_spec = TensorSpec(DType.int64, batch, seqlen)
input_specs = List[InputSpec]()
input_specs.append(input_ids_spec)
input_specs.append(attention_mask_spec)
input_specs.append(token_type_ids_spec)
model = session.load(
Path("../../models/bert-mlm.torchscript"), input_specs=input_specs
)
return model
def read_input() -> String:
"""
コマンドライン引数から入力テキストを読み取ります。
Returns:
String: 入力されたテキスト。
Raises:
Error: 入力が指定されていない場合にエラーをスロー。
"""
USAGE = (
'Usage: ./run.mojo <str> \n\t e.g., ./run.mojo "Paris is the [MASK] of'
' France"'
)
argv = sys.argv()
if len(argv) != 2:
raise Error("\nPlease enter a prompt." + "\n" + USAGE)
return sys.argv()[1]
def main():
"""
プログラムのエントリーポイント。入力テキストを読み取り、モデルを使用して推論を実行します。
コマンドライン引数:
<text>: `[MASK]`トークンを含む入力テキスト。
"""
# Hugging Face Transformersライブラリのインポート
transformers = Python.import_module("transformers")
# 入力テキストの読み取りとモデルのロード
text = read_input()
session = InferenceSession()
model = load_model(session)
# 推論の実行
decoded_result = execute(model, text, transformers)
print("input text: ", text)
print("filled mask: ", text.replace("[MASK]", decoded_result))
Mojoの実行結果
input text: Tokyo is the [MASK] of Japan.
filled mask: Tokyo is the capital of Japan.
magic run bash run.sh 2.65s user 0.46s system 40% cpu 7.763 total
実行環境比較
| 実行環境 |
user時間 |
system時間 |
CPU使用率 | 合計時間 (total) |
|---|---|---|---|---|
| Python | 2.50s | 0.69s | 29% | 10.643s |
| Mojo | 2.65s | 0.46s | 40% | 7.763s |
各列の説明
-
user時間:プログラムがCPUで実行されていた時間 -
system時間:OSがプログラムのために行ったシステムコール（I/O操作やメモリ割り当てなど）の実行時間 - CPU使用率:プログラム実行中に使用されたCPUリソースの割合
-
合計時間 (
total):プログラムの開始から終了までにかかった総時間
Mojoの方が速いですね！とはいえ、思ったより差は小さいかな？それでも、この結果について少し考察します。
まず、system時間はMojoが短く、OSに頼る処理が少なく効率的な設計になってるのかな?
一方で、user時間はMojoが少し長い。これはJITコンパイルの最適化処理が影響している可能性がありますが、全体のパフォーマンスを考えると、十分納得できる範囲かと思います。
CPU使用率は、Pythonが29%なのに対し、Mojoは40%と、リソースを効率よく活用しているのでは？
5. 他言語との比較
ここまで読んで、Juliaでよくね?となった人もいると思うので、MojoがPythonやJuliaとどのように異なり、どの分野で優れているのかについて述べましょう。
Python
Pythonは、そのシンプルな文法と豊富なライブラリが特長で、データ解析や機械学習の分野で広く使用されています。ただし、GIL（グローバルインタプリタロック）による並列処理の制約や、大規模計算でのパフォーマンス不足が課題です。
- 強み: 幅広いエコシステム、豊富なライブラリ、学習のしやすさ
- 課題: 並列処理の制約、実行速度が遅い
Julia
Juliaは科学計算に特化して設計された言語で、高速な数値計算を得意とします。CやFortranに匹敵するパフォーマンスを発揮する一方、エコシステムの規模ではPythonに劣ります。
- 強み: 高速な数値計算、多重ディスパッチによる柔軟性、並列処理のサポート
- 課題: エコシステムが限定的、Pythonとの統合が容易ではない場合がある
Mojo
MojoはPythonの使いやすさを維持しながら、Juliaのような高パフォーマンスを目指しています。AI/MLに特化した設計と、Pythonエコシステムとの高い互換性が特徴です。
- 強み: Pythonライブラリとのシームレスな統合、型安全で効率的な設計
- 課題: 開発初期段階であるため、エコシステムとコミュニティの成熟が必要
結論
Pythonは柔軟性と使いやすさ、Juliaは科学計算や高速処理、そしてMojoはその両者の中間を狙う形で設計されています。
6. まとめ
Mojo🔥いかがだったでしょうか？
現時点では、エコシステムの成熟や実際のプロジェクトでの適用可能性において課題が残るものの、その設計思想は性能と使いやすさの絶妙なバランスを取っていると思います。特にML分野では、高速性とメモリ安全性を兼ね備えたプログラミング言語は、今後必須の能力となる可能性があります。皆さんもぜひ試してみてください！
最後まで読んでいただきありがとうございました!!
Mojoの小ネタ
拡張子が.mojoでも.🔥でもいい
例)attention.🔥