import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
import torch
from threading import Thread
import os  # os 모듈은 혹시 모르니 남겨둡니다.

# --- 1. 모델 로드 (Space의 GPU 활용) ---
# 별도 인증(토큰)이 필요 없는 'kanana-nano' 공개 모델로 변경했습니다.
MODEL_NAME = "kakaocorp/kanana-nano-2.1b-instruct" 

print(f"모델을 로딩 중입니다: {MODEL_NAME} (Space GPU 사용)")

# 이 모델은 공개 모델이므로 토큰(token=...) 파라미터가 필요 없습니다.
tokenizer = AutoTokenizer.from_pretrained(
    MODEL_NAME
)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, 
    torch_dtype=torch.float16,
    device_map="auto"
)
print("모델 로딩 완료!")

# --- 2. 챗봇 응답 함수 (Gradio가 이 함수를 호출) ---
def predict(message, history):
    
    # Kanana의 프롬프트 형식
    history_prompt = ""
    for user_msg, assistant_msg in history:
        history_prompt += f"<bos>user\n{user_msg}\n<eos>assistant\n{assistant_msg}\n"
    
    final_prompt = history_prompt + f"<bos>user\n{message}\n<eos>assistant\n"
    
    inputs = tokenizer(final_prompt, return_tensors="pt").to(model.device)
    
    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
    
    generation_kwargs = dict(
        **inputs,
        streamer=streamer,
        max_new_tokens=1024,
        eos_token_id=tokenizer.eos_token_id,
        pad_token_id=tokenizer.pad_token_id,
        temperature=0.7,
        do_sample=True
    )
    
    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()
    
    generated_text = ""
    for new_text in streamer:
        generated_text += new_text
        yield generated_text

# --- 3. Gradio 챗봇 UI 생성 ---
chatbot_ui = gr.ChatInterface(
    fn=predict,
    title="Kanana-Nano 챗봇 테스트 🤖",
    description=f"{MODEL_NAME} 모델을 테스트합니다.",
    theme="soft",
    examples=[["한국의 수도는 어디야?"], ["AI에 대해 3줄로 요약해줘."]]
)

# --- 4. 앱 실행 ---
chatbot_ui.launch()