import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_name = "dasomaru/gemma-3-4bit-it-demo"

# 🚀 tokenizer는 CPU에서도 미리 불러올 수 있음
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
# 🚀 model은 CPU로만 먼저 올림 (GPU 아직 없음)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,  # 4bit model이니까
    trust_remote_code=True,
)

@spaces.GPU(duration=300)
def generate_response(prompt):
    # 모델 및 토크나이저 로딩은 함수 내부에서 수행
    tokenizer = AutoTokenizer.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
    model = AutoModelForCausalLM.from_pretrained("dasomaru/gemma-3-4bit-it-demo")
    model.to("cuda")

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
    outputs = model.generate(**inputs, max_new_tokens=512, temperature=0.7,
    top_p=0.9,
    top_k=50,
    do_sample=True,)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(fn=generate_response, inputs="text", outputs="text")
demo.launch()



# zero = torch.Tensor([0]).cuda()
# print(zero.device) # <-- 'cpu' 🤔

# @spaces.GPU
# def greet(n):
#     print(zero.device) # <-- 'cuda:0' 🤗
#     return f"Hello {zero + n} Tensor"

# demo = gr.Interface(fn=greet, inputs=gr.Number(), outputs=gr.Text())
# demo.launch()