from transformers import AutoTokenizer
from transformers import AutoModelForCausalLM
import torch

CHECKPOINT = "HuggingFaceTB/SmolLM2-1.7B-Instruct"

_model = None
_tokenizer = None

try:
    import spaces
    IS_SPACES = True
except ImportError:
    IS_SPACES = False


def _load_model(device):
    global _model, _tokenizer
    if _model is None:
        _tokenizer = AutoTokenizer.from_pretrained(CHECKPOINT)
        _model = AutoModelForCausalLM.from_pretrained(CHECKPOINT).to(device)


def _generate(prompt, device):
    _load_model(device)

    messages = [{"role": "user", "content": prompt}]
    chat = _tokenizer.apply_chat_template(messages, tokenize=False)
    inputs = _tokenizer(chat, return_tensors="pt").to(device)

    input_len = inputs["input_ids"].shape[1]

    outputs = _model.generate(
        **inputs,
        max_new_tokens=2048,
        temperature=0.2,
        top_p=0.9,
        do_sample=True,
        pad_token_id=_tokenizer.eos_token_id
    )

    generated_tokens = outputs[0][input_len:]
    response = _tokenizer.decode(generated_tokens, skip_special_tokens=True).strip()
    print(response)
    return response


if IS_SPACES:
    @spaces.GPU
    def generate(prompt):
        return _generate(prompt, "cuda")
else:
    def generate(prompt):
        device = "cuda" if torch.cuda.is_available() else "cpu"
        return _generate(prompt, device)