# -*- coding: utf-8 -*-
"""tinyllama.ipynb

Automatically generated by Colab.

Original file is located at
    https://colab.research.google.com/drive/1hptURmKe_W0-VjcQcLweGBPsjuKJXCV4
"""

# Commented out IPython magic to ensure Python compatibility.
# %%bash
# set -euxo pipefail
# 
# pip install -U pip
# # Clean out conflicting installs
# pip uninstall -y torch torchvision torchaudio bitsandbytes accelerate || true
# 
# # Install CUDA-enabled PyTorch for Colab (CUDA 12.1). Fall back to cu118 if needed.
# pip install --index-url https://download.pytorch.org/whl/cu121 torch torchvision torchaudio \
#   || pip install --index-url https://download.pytorch.org/whl/cu118 torch torchvision torchaudio
# 
# # Install bitsandbytes (>=0.43.1 to avoid old CUDA-only import paths) + transformers + sentencepiece
# pip install "bitsandbytes>=0.43.1" "transformers>=4.44.0" sentencepiece
#

import torch, bitsandbytes as bnb
print("PyTorch:", torch.__version__)
print("CUDA available:", torch.cuda.is_available())
print("bitsandbytes:", bnb.__version__)

from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline

model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,   # FP16 on GPU
    device_map=None
).to("cuda")                     # or "cpu" if no GPU

chat = pipeline("text-generation", model=model, tokenizer=tokenizer, device=0)


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

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

print(response)