File size: 3,016 Bytes
c39d6d1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 | import torch
import numpy as np
import os
from config import ZUCOConfig
from model import build_model
from transformers import AutoTokenizer
def test_model_basic():
print("\n" + "="*80)
print("bai-Mind-8-v1")
print("="*80 + "\n")
model_path = r"model/path"
config = ZUCOConfig()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"[INFO] Device: {device}")
print(f"\n[1/4] Creating Model...")
model = build_model(config)
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f" Total Params: {total_params:,}")
print(f" Trainable Params: {trainable_params:,}")
print(f"\n[2/4] Model checkpoint loading...")
print(f" Path: {model_path}")
if os.path.exists(model_path):
checkpoint = torch.load(model_path, map_location=device)
model.load_state_dict(checkpoint['model_state_dict'])
print(f" ✓ Model loaded successfully!")
if 'epoch' in checkpoint:
print(f" Epoch: {checkpoint['epoch']}")
if 'val_loss' in checkpoint:
print(f" Validation Loss: {checkpoint['val_loss']:.4f}")
else:
print(f" ✗ Model path is not found!")
print(f" Using randomize model")
model.to(device)
model.eval()
# Tokenizer
print(f"\n[3/4] Tokenizer loading...")
tokenizer = AutoTokenizer.from_pretrained(config.language_model_name)
print(f" ✓ Tokenizer is ready")
print(f"\n[4/4] Test...")
batch_size = 2
eeg_seq_len = 50
eeg = torch.randn(batch_size, eeg_seq_len, config.num_eeg_channels).to(device)
eeg_attention_mask = torch.ones(batch_size, eeg_seq_len).to(device)
print(f" Test EEG shape: {eeg.shape}")
# Inference
with torch.no_grad():
generated_ids = model.generate(
eeg=eeg,
eeg_attention_mask=eeg_attention_mask,
max_length=50,
num_beams=2,
early_stopping=True
)
# Decode
predictions = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
print(f" ✓ Model is running succesfully!")
print(f"\n Sample predictions:")
for i, pred in enumerate(predictions):
print(f" [{i+1}] {pred}")
print("\n" + "="*80)
print("TEST RESULT: SUCCESS ✓")
print("="*80 + "\n")
return {
'status': 'success',
'total_params': total_params,
'trainable_params': trainable_params,
'device': str(device),
'sample_predictions': predictions
}
if __name__ == "__main__":
try:
results = test_model_basic()
print("[OK] Model test completed!")
except Exception as e:
print(f"\n[ERROR]: {e}")
import traceback
traceback.print_exc()
|