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

model_name = "kakaocorp/kanana-nano-2.1b-instruct"

# 모델과 토크나이저를 CPU 환경에서 로드합니다.
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float32,  # CPU에서는 bfloat16 지원이 제한될 수 있으므로 float32 사용 권장
    trust_remote_code=True,
)
# CPU만 사용하므로 .to("cuda") 부분은 생략합니다.
tokenizer = AutoTokenizer.from_pretrained(model_name)

def generate_response(prompt):
    messages = [
        {"role": "system", "content": "You are a helpful AI assistant developed by Kakao."},
        {"role": "user", "content": prompt}
    ]
    input_ids = tokenizer.apply_chat_template(
        messages,
        tokenize=True,
        add_generation_prompt=True,
        return_tensors="pt"
    )
    
    model.eval()
    with torch.no_grad():
        output = model.generate(
            input_ids,
            max_new_tokens=72,
            do_sample=False,
        )
    return tokenizer.decode(output[0], skip_special_tokens=True)

with gr.Blocks() as demo:
    with gr.Tab("About"):
        gr.Markdown("# Inference Provider")
        gr.Markdown("이 Space는 kakaocorp/kanana-nano-2.1b-instruct 모델을 CPU에서 추론합니다.")
    
    with gr.Tab("Generate"):
        prompt_input = gr.Textbox(
            label="Prompt 입력",
            placeholder="여기에 프롬프트를 입력하세요...",
            lines=5
        )
        generate_btn = gr.Button("생성")
        output_text = gr.Textbox(
            label="모델 출력",
            lines=10
        )
        generate_btn.click(fn=generate_response, inputs=prompt_input, outputs=output_text)

demo.launch()