kurogane/

Nemotron-H-micro-test03

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "kurogane/Nemotron-H-micro-test03"
dir_cache = r"/media/kurogane/backup/cache"

# Load tokenizer and model
tokenizer = AutoTokenizer.from_pretrained(
    model_id,
    trust_remote_code=True,
    cache_dir=dir_cache,
)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    trust_remote_code=True,
    device_map="auto",
    cache_dir=dir_cache,
)

inputs = tokenizer(
    "こんにちは", 
    return_tensors="pt",
).to(model.device)

outputs = model.generate(**inputs)
print(
    tokenizer.batch_decode(
        outputs, 
        skip_special_tokens=True,
    )
)