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)