MarkBase Admin
ac75faa0cc
CI / build-and-test (push) Has been cancelled
- E4B-MarkBase model (42 layers, 4.4GB) loaded successfully - All Phase 1-6 tests passed (model loading, forward pass, vision/audio towers, token generation, performance) - All stress tests passed (5/5 in 127.6s) - Concurrent inference - Memory stress (67.5 tok/s, 0 NaN) - Continuous generation - Batch processing - Long-running stability - Swift Metal inference engine with multimodal support
84 lines
2.6 KiB
Python
84 lines
2.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Compare logits from original E4B model vs our Swift implementation.
|
|
Uses transformers library since MLX doesn't support Gemma-4 yet.
|
|
"""
|
|
|
|
import torch
|
|
from transformers import AutoModelForCausalLM, AutoTokenizer
|
|
import json
|
|
import numpy as np
|
|
|
|
# Load original E4B model from HuggingFace
|
|
model_path = "google/gemma-4-4b-it"
|
|
|
|
print("Loading original E4B model from HuggingFace...")
|
|
print("This may take a moment to download...")
|
|
|
|
try:
|
|
tokenizer = AutoTokenizer.from_pretrained(model_path)
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
model_path,
|
|
device_map="auto",
|
|
torch_dtype=torch.bfloat16
|
|
)
|
|
print("Model loaded successfully!")
|
|
except Exception as e:
|
|
print(f"Error loading model: {e}")
|
|
print("\nFalling back to checking config from converted model...")
|
|
|
|
# Check config from converted model
|
|
conv_path = "/Users/accusys/MarkBase12B/models/E4B-MarkBase/config.json"
|
|
with open(conv_path) as f:
|
|
cfg = json.load(f)
|
|
|
|
print("\nConverted model config:")
|
|
for key in ["hidden_size", "num_hidden_layers", "vocab_size",
|
|
"num_attention_heads", "num_key_value_heads",
|
|
"hidden_size_per_layer_input"]:
|
|
print(f" {key}: {cfg.get('text_config', {}).get(key, cfg.get(key))}")
|
|
|
|
exit(1)
|
|
|
|
# Get BOS token
|
|
bos_token_id = tokenizer.bos_token_id if tokenizer.bos_token_id else 2
|
|
print(f"\nBOS token ID: {bos_token_id}")
|
|
|
|
# Run forward pass for position 0
|
|
input_ids = torch.tensor([[bos_token_id]])
|
|
print(f"Input IDs: {input_ids}")
|
|
|
|
with torch.no_grad():
|
|
outputs = model(input_ids)
|
|
logits = outputs.logits[0, 0].float().numpy() # [vocab_size]
|
|
|
|
print(f"\nLogits shape: {logits.shape}")
|
|
print(f"Logits range: min={logits.min():.2f}, max={logits.max():.2f}")
|
|
|
|
# Get top 10 tokens
|
|
top_indices = np.argsort(logits)[-10:][::-1]
|
|
print("\nTop 10 tokens (position 0, BOS):")
|
|
for i, idx in enumerate(top_indices):
|
|
token = tokenizer.decode([idx])
|
|
print(f" {i+1}. token {idx} '{token}': {logits[idx]:.2f}")
|
|
|
|
# Test with actual prompt
|
|
prompt = "<start_of_turn>user\nThe capital of France is<end_of_turn>\n<start_of_turn>model\n"
|
|
print(f"\nTesting prompt: '{prompt[:50]}...'")
|
|
|
|
# Tokenize
|
|
inputs = tokenizer(prompt, return_tensors="pt")
|
|
print(f"Tokens: {inputs['input_ids'][0][:20].tolist()}...")
|
|
|
|
# Generate
|
|
print("\nGenerating response...")
|
|
with torch.no_grad():
|
|
outputs = model.generate(
|
|
**inputs,
|
|
max_new_tokens=30,
|
|
do_sample=True,
|
|
temperature=1.0,
|
|
top_k=40
|
|
)
|
|
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
|
|
print(f"Response: '{response}'") |