ananjayram/hana

TinyMoELM - Orange MoE

This is a custom Mixture of Experts (MoE) Language Model.

How to Load and Use

Since this model uses a custom architecture, you can initialize the architecture and load the safetensors directly:

import torch
from transformers import AutoTokenizer
from safetensors.torch import load_file

# 1. Define or import Config and TinyMoELM class from original code
# (Make sure Config and TinyMoELM classes are defined in your script)

# 2. Initialize Config and Model
config_data = {'vocab_size': 50257, 'block_size': 512, 'd_model': 512, 'n_layers': 8, 'n_heads': 8, 'use_moe': True, 'num_experts': 8, 'top_k': 2, 'expert_ffn_dim': 768}
cfg = Config(**config_data)
model = TinyMoELM(cfg)

# 3. Load Safetensors weights directly
from huggingface_hub import hf_hub_download
weights_path = hf_hub_download(repo_id="ananjayram/hana", filename="model.safetensors")
state_dict = load_file(weights_path)
model.load_state_dict(state_dict)
model.eval()

# 4. Tokenize and Generate
tokenizer = AutoTokenizer.from_pretrained("gpt2")
prompt = "Once upon a time"
inputs = tokenizer(prompt, return_tensors="pt")["input_ids"]
with torch.no_grad():
    logits, _, _ = model(inputs)