import spaces
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel
import torch
import os

token = os.environ.get("HF_TOKEN")
BASE_MODEL = "unsloth/Llama-3.2-3B-bnb-4bit"
ADAPTER = "Ganesh3108/Second-model"

# only load tokenizer at module level — it's not a CUDA op
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=token)

model = None  # lazy-loaded

def load_model():
    global model
    if model is None:
        base_model = AutoModelForCausalLM.from_pretrained(
            BASE_MODEL,
            torch_dtype=torch.float16,
            device_map={"": 0},   # or "auto"
            token=token,
        )
        model = PeftModel.from_pretrained(base_model, ADAPTER)
        model.eval()
    return model

@spaces.GPU
def chat(message, history):
    m = load_model()  # loads on first call, when GPU is actually attached
    inputs = tokenizer(message, return_tensors="pt").to("cuda")
    with torch.no_grad():
        outputs = m.generate(
            **inputs,
            max_new_tokens=150,
            do_sample=True,
            temperature=1.5,
            top_p=0.9,
            repetition_penalty=1.2,
            no_repeat_ngram_size=3,
            eos_token_id=tokenizer.eos_token_id,
            pad_token_id=tokenizer.eos_token_id,
        )
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response[len(message):].strip()

demo = gr.ChatInterface(chat, title="My Bro")
demo.launch()