File size: 1,108 Bytes
62e57e7 | 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 | from pathlib import Path
import sys
# Exact RUMI public package directory
PACKAGE_ROOT = Path(__file__).resolve().parents[1]
# Put the public package directory first in Python's import path
sys.path.insert(0, str(PACKAGE_ROOT))
print("Package root:", PACKAGE_ROOT)
print("Checking RUMI package...")
from rumi import RUMIModel
print("RUMI package imported successfully.")
print()
print("Loading RUMI-8M...")
model = RUMIModel.from_pretrained(PACKAGE_ROOT)
print()
print("RUMI loaded successfully.")
print("Model:", model.model_name)
print("Parameters:", f"{model.parameter_count:,}")
print("Vocabulary:", model.vocab_size)
print("Context length:", model.context_length)
print("Embedding dimension:", model.embedding_dim)
print("Layers:", model.num_layers)
print("Attention heads:", model.num_heads)
print("FFN dimension:", model.feed_forward_dim)
print("Device:", model.device)
print()
print("Generating text...")
output = model.generate(
"Artificial intelligence is",
max_new_tokens=30,
temperature=0.8,
top_k=50,
do_sample=True,
)
print()
print("Generated text:")
print(output) |