# ✱ ONE-CELL DEMO: prompt TinyLlama-1.1B-Chat and print its reply ✱
#
# ▶ Requires ≈2 GB RAM if you let Transformers load it in 4-bit.
# ▶ Works on laptop GPUs or free-tier Colab (T4 / K80).
# -----------------------------------------------------------------

# 1.  Install minimal deps (≈ 2-3 min the first time)
# !pip install -q transformers accelerate bitsandbytes

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"   # open-weights on HF Hub

# 2.  Load tokenizer + 4-bit quantised model onto GPU/CPU automatically
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    load_in_4bit=True,       # bitsandbytes quantisation → tiny VRAM
    device_map="auto",       # sends weights to first CUDA device if available
    bnb_4bit_compute_dtype=torch.float16,   # <─ compute in FP16
    bnb_4bit_quant_type="nf4",              # (default is fine, just explicit)
    bnb_4bit_use_double_quant=True          # optional, tiny accuracy boost
)

# 3.  Quick generation pipeline
chat = pipeline("text-generation", model=model, tokenizer=tokenizer)

prompt = (
    "### User: In one sentence, explain why the sky is blue.\n"
    "### Assistant:"
)

# prompt = (
#     "### User: Why is a raven like a writing desk?\n"
#     "### Assistant:"
# )

response = chat(
    prompt,
    max_new_tokens=1024,
    temperature=0.3,#0.7,
    do_sample=True,
)[0]["generated_text"]

print(response)

The `load_in_4bit` and `load_in_8bit` arguments are deprecated and will be removed in the future versions. Please, pass a `BitsAndBytesConfig` object in `quantization_config` argument instead.
Device set to use cuda:0

### User: In one sentence, explain why the sky is blue.
### Assistant: The sky is blue because the colors of the colors of the sky are arranged in a way that creates the appearance of a blue sky. The blue is created by the combination of blue, green, and red colors. The blue color is created by the blue and green colors, and the red color is created by the red and orange colors. The colors are arranged in a way that creates a blue sky, which is the result of the interplay of these colors.