# 1. Install dependencies
# pip install transformers torch huggingface_hub

import sys
import os
import torch
from transformers import AutoTokenizer
from huggingface_hub import snapshot_download

# 2. Download the repository (Code + Weights)
repo_path = snapshot_download(repo_id="Adx19/gemma-3-270m-tinystories")

# 3. Add the downloaded folder to Python path so we can import 'src'
sys.path.append(repo_path)

# 4. Import Custom Model
from src.model import Gemma3Model
from src.config import GEMMA3_CONFIG_270M

# 5. Load Model & Weights
device = "cuda" if torch.cuda.is_available() else "cpu"
model = Gemma3Model(GEMMA3_CONFIG_270M).to(device)

weights_path = os.path.join(repo_path, "pytorch_model.bin")
model.load_state_dict(torch.load(weights_path, map_location=device))
model.eval()

# 6. Load Tokenizer
tokenizer = AutoTokenizer.from_pretrained("Adx19/gemma-3-270m-tinystories")

# 7. Generate
input_text = "Once upon a time"
input_ids = tokenizer(input_text, return_tensors="pt")["input_ids"].to(device)

with torch.no_grad():
    output = model.generate(input_ids, max_new_tokens=50)

print(tokenizer.decode(output[0], skip_special_tokens=True))