33 lines
754 B
Python
33 lines
754 B
Python
import os
|
|
import sys
|
|
|
|
ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", ".."))
|
|
sys.path.insert(0, ROOT)
|
|
|
|
import torch
|
|
import transformers
|
|
import datasets
|
|
|
|
from model.model_minimind import MiniMindConfig, MiniMindForCausalLM
|
|
|
|
|
|
def main():
|
|
config = MiniMindConfig(
|
|
hidden_size=128,
|
|
num_hidden_layers=1,
|
|
num_attention_heads=4,
|
|
num_key_value_heads=2,
|
|
vocab_size=6400,
|
|
)
|
|
model = MiniMindForCausalLM(config)
|
|
params = sum(p.numel() for p in model.parameters())
|
|
|
|
print(f"torch={torch.__version__}")
|
|
print(f"transformers={transformers.__version__}")
|
|
print(f"datasets={datasets.__version__}")
|
|
print(f"tiny_minimind_params={params}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|