KoGPT-Chat / app.py
Yuchan5386's picture
Update app.py
8ac5e40 verified
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)