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

try:
    model_name = "gpt2"

    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)

    chat_history = [{"role": "system", "content": "You are a helpful assistant."}]

    def chat(user_input):
        # 사용자 입력을 채팅 기록에 추가
        chat_history.append({"role": "user", "content": user_input})
        
        # 모델에 채팅 기록 전달
        inputs = tokenizer([message['content'] for message in chat_history], return_tensors="pt", padding=True).to("cuda" if torch.cuda.is_available() else "cpu")
        
        # 모델로 응답 생성
        outputs = model.generate(**inputs, max_length=200)
        
        # 생성된 응답
        bot_reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
        
        # 봇 응답을 채팅 기록에 추가
        chat_history.append({"role": "assistant", "content": bot_reply})
        
        return bot_reply

    demo = gr.ChatInterface(fn=chat, type='messages', title='이우진의 챗봇')
    demo.launch(share=True)
except Exception as e:
    print(f"Error occurred: {e}")