# 모델 로딩
import torch
from peft import PeftConfig, PeftModel
from transformers import AutoModelForCausalLM, AutoTokenizer

device = "cuda" if torch.cuda.is_available() else "cpu"

base_model_name = "facebook/opt-350m"
adapter_model_name = 'msy127/opt-350m-aihubqa-130-dpo-adapter'

model = AutoModelForCausalLM.from_pretrained(base_model_name)
model = PeftModel.from_pretrained(model, adapter_model_name).to(device)
tokenizer = AutoTokenizer.from_pretrained(adapter_model_name)

# 대화 누적 함수 (history) - prompt 자리에 history가 들어감 -> dialoGPT는 모델 집어넣기 전에 인코딩을 했었는데 OPENAI는 인코딩을 안한다.

def predict(input, history):
    history.append({"role": "user", "content": input})

    # 일반모델
    prompt = f"An AI tool that looks at the context and question separated by triple backquotes, finds the answer corresponding to the question in the context, and answers clearly.\n### Input: ```{input}```\n ### Output: "
    inputs = tokenizer.encode(prompt, return_tensors="pt").to(device)
    outputs = model.generate(input_ids=inputs, max_length=256)
    generated_text = tokenizer.decode(outputs[0])
    start_idx = len(prompt) + len('</s>')
    stop_first_idx = generated_text.find("### Input:") # 첫 번째 "### Input:"을 찾습니다.
    stop_idx = generated_text.find("### Input:", stop_first_idx + 1) # 첫 번째 "### Input:" 이후의 문자열에서 다시 "### Input:"을 찾습니다.
    # print(start_idx , stop_idx)
    # print(generated_text)
    if stop_idx != -1:
        response = generated_text[start_idx:stop_idx] # prompt 뒤에 있는 새롭게 생성된 텍스트만 ("### Input:" 전까지) 가져옵니다.

    # 누적
    history.append({"role": "assistant", "content": response})
    # messages = [(history[i]["content"], history[i+1]["content"]) for i in range(1, len(history), 2)]
    messages = [(history[i]["content"], history[i+1]["content"]) for i in range(0, len(history) - 1, 2)]

    return messages, history


# Gradio 인터페이스 설정
import gradio as gr
with gr.Blocks() as demo:
    chatbot = gr.Chatbot(label="ChatBot")

    state = gr.State([
        {"role": "system", "content": "당신은 친절한 인공지능 챗봇입니다. 입력에 대해 짧고 간결하고 친절하게 대답해주세요."}])

    with gr.Row():
        txt = gr.Textbox(show_label=False, placeholder="챗봇에게 아무거나 물어보세요").style(container=False)
        # txt.submit(predict, [txt, state], [chatbot, state])

    txt.submit(predict, [txt, state], [chatbot, state])

# demo.launch(debug=True, share=True)
demo.launch()



# from PIL import Image
# import gradio as gr
# interface = gr.Interface(
#     fn=classify_image,
#     inputs=gr.components.Image(type="pil", label="Upload an Image"),
#     outputs="text",
#     live=True
# )
# interface.launch()