import os
import gradio as gr
from transformers import GPT2LMHeadModel, PreTrainedTokenizerFast

MODEL_NAME = "skt/kogpt2-base-v2"  # 사용하고자 하는 KoGPT2 모델
API_TOKEN = os.getenv("HUGGINGFACE_API_TOKEN")  # 환경 변수에서 토큰 불러오기

# PreTrainedTokenizerFast로 변경하여 오류 해결
tokenizer = PreTrainedTokenizerFast.from_pretrained(MODEL_NAME, use_auth_token=API_TOKEN)
model = GPT2LMHeadModel.from_pretrained(MODEL_NAME, use_auth_token=API_TOKEN)

def kogpt2_chatbot(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    output = model.generate(input_ids, max_length=100, temperature=0.7, top_p=0.9, repetition_penalty=1.2)
    response_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return response_text

# Gradio 인터페이스 설정
iface = gr.Interface(fn=kogpt2_chatbot, inputs="text", outputs="text", title="KoGPT2 Chatbot")

# 외부에서도 접근할 수 있는 공개 링크를 생성
iface.launch(share=True)