# Demo: NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16 – Textgenerierung mit transformers
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

REPO = "nvidia/NVIDIA-Nemotron-Labs-Dory1.0-4B-BF16"
PARAMS_B = 4.0

kwargs = dict(torch_dtype="auto", device_map="auto", trust_remote_code=True)
if torch.cuda.is_available() and PARAMS_B >= 7:
    try:
        from transformers import BitsAndBytesConfig
        kwargs["quantization_config"] = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.float16)
        print("4-bit-Quantisierung aktiv (spart VRAM)")
    except Exception as e:
        print("bitsandbytes nicht verfügbar, lade ohne Quantisierung:", e)

tok = AutoTokenizer.from_pretrained(REPO, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(REPO, **kwargs)

messages = [{"role": "user", "content": "Erkläre in drei Sätzen, was ein Large Language Model ist."}]
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt", return_dict=True)
inputs = {k: v.to(model.device) for k, v in inputs.items()}
out = model.generate(**inputs, max_new_tokens=400, do_sample=True, temperature=0.7)
print(tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True))