app.py

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  # 이 함수 실행될 때 GPU 할당됨!
def chat(user_input):
    model.to("cuda")  # 함수 안에서 GPU로 이동!
    
    messages = [{
        "role": "user",
        "content": [{"type": "text", "text": user_input}]
    }]
    prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False)

    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")

    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=256,
            temperature=1.0,
            top_p=0.95,
            top_k=64,
            do_sample=True,
        )

    output_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return output_text.split(user_input)[-1].strip()

# Gradio 인터페이스 설정
demo = gr.Interface(
    fn=chat,
    inputs=gr.Textbox(lines=2, placeholder="Type your message..."),
    outputs=gr.Textbox(lines=10),
    title="🧠 Gemma-3 4bit (ZeroGPU)",
    description="This Space uses the ZeroGPU feature. First request might take a few seconds!"
)

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()