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

MODEL_NAME = "lcw99/ko-dialoGPT-korean-chit-chat"

# 토크나이저 및 모델 로드
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForCausalLM.from_pretrained(MODEL_NAME)

# 챗봇 응답 함수
def chat_with_ai(history, message):
    input_text = message + tokenizer.eos_token
    input_ids = tokenizer.encode(input_text, return_tensors="pt")

    response_ids = model.generate(input_ids, max_length=100, pad_token_id=tokenizer.eos_token_id)
    response_text = tokenizer.decode(response_ids[:, input_ids.shape[-1]:][0], skip_special_tokens=True)

    history.append((message, response_text))  # Gradio Chatbot 형식 유지

    return history, ""  # 입력창 비우기

# Gradio 인터페이스 설정
with gr.Blocks() as demo:
    gr.Markdown("## 🗨️ Ko-DialoGPT Chatbot")
    
    chatbot = gr.Chatbot(label="Ko-DialoGPT Chatbot")
    message = gr.Textbox(label="입력 메시지", placeholder="메시지를 입력하세요...")
    clear_btn = gr.Button("초기화")

    # 메시지 입력 시 채팅 기록 업데이트 및 입력창 초기화
    message.submit(chat_with_ai, [chatbot, message], [chatbot, message])

    # 초기화 버튼 클릭 시 채팅 기록 삭제
    clear_btn.click(lambda: [], [], chatbot)

# ✅ 서버 포트 및 주소 추가
demo.launch(server_name="0.0.0.0", server_port=7860)