app.py

import gradio as gr
from huggingface_hub import InferenceClient
import openai
import anthropic
from typing import Optional

#############################
# [기본코드] - 수정/삭제 불가
#############################

# 제거할 모델들을 MODELS 사전에서 제외
MODELS = {
    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
    "Meta Llama 3.1 8B": "meta-llama/Meta-Llama-3.1-8B-Instruct",
    "Meta-Llama 3.1 70B-Instruct": "meta-llama/Meta-Llama-3.1-70B-Instruct",
    "Microsoft": "microsoft/Phi-3-mini-4k-instruct",
    "Mixtral 8x7B": "mistralai/Mistral-7B-Instruct-v0.3",
    "Mixtral Nous-Hermes": "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
    "Aya-23-35B": "CohereForAI/aya-23-35B"
}

# Cohere Command R+ 모델 ID 정의
COHERE_MODEL = "CohereForAI/c4ai-command-r-plus-08-2024"

def get_client(model_name, hf_token):
    """
    모델 이름에 맞춰 InferenceClient 생성.
    hf_token을 UI에서 입력받은 값으로 사용하도록 변경.
    """
    if not hf_token:
        raise ValueError("HuggingFace API 토큰이 필요합니다.")

    if model_name in MODELS:
        model_id = MODELS[model_name]
    elif model_name == "Cohere Command R+":
        model_id = COHERE_MODEL
    else:
        raise ValueError("유효하지 않은 모델 이름입니다.")
    return InferenceClient(model_id, token=hf_token)


def respond_hf_qna(
    question: str,
    model_name: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    system_message: str,
    hf_token: str
):
    """
    HuggingFace 모델(Zephyr 등)에 대해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
    """
    try:
        client = get_client(model_name, hf_token)
    except ValueError as e:
        return f"오류: {str(e)}"

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": question}
    ]

    try:
        response = client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            stream=False,
        )
        assistant_message = response.choices[0].message.content
        return assistant_message

    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"


def respond_cohere_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    hf_token: str
):
    """
    Cohere Command R+ 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
    """
    model_name = "Cohere Command R+"
    try:
        client = get_client(model_name, hf_token)
    except ValueError as e:
        return f"오류: {str(e)}"

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": question}
    ]

    try:
        response_full = client.chat_completion(
            messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response_full.choices[0].message.content
        return assistant_message
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"


def respond_chatgpt_qna(
    question: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    openai_token: str
):
    """
    ChatGPT(OpenAI) 모델을 이용해 한 번의 질문(question)에 대한 답변을 반환하는 함수.
    """
    if not openai_token:
        return "OpenAI API 토큰이 필요합니다."

    openai.api_key = openai_token

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": question}
    ]

    try:
        response = openai.ChatCompletion.create(
            model="gpt-4o-mini",  # 필요한 경우 변경
            messages=messages,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
        )
        assistant_message = response.choices[0].message['content']
        return assistant_message
    except Exception as e:
        return f"오류가 발생했습니다: {str(e)}"


#############################
# [기본코드] UI 부분 - 수정/삭제 불가
#############################

with gr.Blocks() as demo:
    gr.Markdown("# Prompting AI - 일반 문답형 데모")
    gr.Markdown("언어모델별 문답형 테스트 데모입니다. 한 번에 한 질문씩만 주고받습니다.")

    # 한 줄에 세 토큰 텍스트박스 배치
    with gr.Row():
        hf_token_box = gr.Textbox(
            label="HuggingFace 토큰 (비공개)",
            type="password",
            placeholder="HuggingFace API 토큰을 입력하세요..."
        )
        openai_token_box = gr.Textbox(
            label="OpenAI 토큰 (비공개)",
            type="password",
            placeholder="OpenAI API 토큰을 입력하세요..."
        )
        claude_token_box = gr.Textbox(
            label="Claude 토큰 (비공개)",
            type="password",
            placeholder="Claude API 토큰을 입력하세요...",
            show_copy_button=False
        )

    #################
    # 일반 모델 탭
    #################
    with gr.Tab("일반 모델"):
        # 모델명 선택
        model_name = gr.Radio(
            choices=list(MODELS.keys()),
            label="Language Model (HuggingFace)",
            value="Zephyr 7B Beta"
        )

        # 입력1~5 (한 줄)
        with gr.Row():
            input1 = gr.Textbox(label="입력1", lines=1)
            input2 = gr.Textbox(label="입력2", lines=1)
            input3 = gr.Textbox(label="입력3", lines=1)
            input4 = gr.Textbox(label="입력4", lines=1)
            input5 = gr.Textbox(label="입력5", lines=1)

        # 답변
        answer_output = gr.Textbox(label="답변", lines=5, interactive=False)

        # 고급 설정을 답변 아래에
        with gr.Accordion("고급 설정 (일반 모델)", open=False):
            max_tokens = gr.Slider(minimum=0, maximum=2000, value=500, step=100, label="Max Tokens")
            temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")
            system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 최고의 비서이다.
내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )

        submit_button = gr.Button("전송")

        # 다섯 입력칸을 합쳐서 question으로 만든 뒤 응답
        def merge_and_call_hf(i1, i2, i3, i4, i5, m_name, mt, temp, top_p_, sys_msg, hf_token):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_hf_qna(
                question=question,
                model_name=m_name,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_,
                system_message=sys_msg,
                hf_token=hf_token
            )

        submit_button.click(
            fn=merge_and_call_hf,
            inputs=[
                input1, input2, input3, input4, input5,  # 입력1~5
                model_name,
                max_tokens,
                temperature,
                top_p,
                system_message,
                hf_token_box
            ],
            outputs=answer_output
        )

    #################
    # Cohere Command R+ 탭
    #################
    with gr.Tab("Cohere Command R+"):
        # 입력1~5 (한 줄)
        with gr.Row():
            cohere_input1 = gr.Textbox(label="입력1", lines=1)
            cohere_input2 = gr.Textbox(label="입력2", lines=1)
            cohere_input3 = gr.Textbox(label="입력3", lines=1)
            cohere_input4 = gr.Textbox(label="입력4", lines=1)
            cohere_input5 = gr.Textbox(label="입력5", lines=1)

        # 답변
        cohere_answer_output = gr.Textbox(label="답변", lines=5, interactive=False)

        with gr.Accordion("고급 설정 (Cohere)", open=False):
            cohere_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 최고의 비서이다.
내가 요구하는것들을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            cohere_max_tokens = gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max Tokens")
            cohere_temperature = gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature")
            cohere_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")

        cohere_submit_button = gr.Button("전송")

        def merge_and_call_cohere(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, hf_token):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_cohere_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_,
                hf_token=hf_token
            )

        cohere_submit_button.click(
            fn=merge_and_call_cohere,
            inputs=[
                cohere_input1, cohere_input2, cohere_input3, cohere_input4, cohere_input5,
                cohere_system_message,
                cohere_max_tokens,
                cohere_temperature,
                cohere_top_p,
                hf_token_box
            ],
            outputs=cohere_answer_output
        )

    #################
    # ChatGPT 탭
    #################
    with gr.Tab("gpt-4o-mini"):
        # 입력1~5 (한 줄)
        with gr.Row():
            chatgpt_input1 = gr.Textbox(label="입력1", lines=1)
            chatgpt_input2 = gr.Textbox(label="입력2", lines=1)
            chatgpt_input3 = gr.Textbox(label="입력3", lines=1)
            chatgpt_input4 = gr.Textbox(label="입력4", lines=1)
            chatgpt_input5 = gr.Textbox(label="입력5", lines=1)

        # 답변
        chatgpt_answer_output = gr.Textbox(label="답변", lines=5, interactive=False)

        with gr.Accordion("고급 설정 (ChatGPT)", open=False):
            chatgpt_system_message = gr.Textbox(
                value="""반드시 한글로 답변할 것.
너는 ChatGPT, OpenAI에서 개발한 언어 모델이다.
내가 요구하는 것을 최대한 자세하고 정확하게 답변하라.
""",
                label="System Message",
                lines=3
            )
            chatgpt_max_tokens = gr.Slider(minimum=1, maximum=4096, value=1024, step=1, label="Max Tokens")
            chatgpt_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.05, label="Temperature")
            chatgpt_top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-P")

        chatgpt_submit_button = gr.Button("전송")

        def merge_and_call_chatgpt(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, openai_token):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_chatgpt_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_,
                openai_token=openai_token
            )

        chatgpt_submit_button.click(
            fn=merge_and_call_chatgpt,
            inputs=[
                chatgpt_input1, chatgpt_input2, chatgpt_input3, chatgpt_input4, chatgpt_input5,
                chatgpt_system_message,
                chatgpt_max_tokens,
                chatgpt_temperature,
                chatgpt_top_p,
                openai_token_box
            ],
            outputs=chatgpt_answer_output
        )

    #################################################
    # [클로드 플레이그라운드] - 개선된 코드
    #################################################

    def validate_claude_token(token: str) -> bool:
        """Claude API 토큰 검증"""
        return bool(token and len(token.strip()) >= 10)

    def respond_claude_qna(
        question: str,
        system_message: str,
        max_tokens: int,
        temperature: float,
        top_p: float,
        claude_api_key: str
    ) -> str:
        """
        Claude API를 사용한 개선된 응답 생성 함수
        """
        if not validate_claude_token(claude_api_key):
            return "유효한 Claude API 토큰이 필요합니다."
    
        try:
            client = anthropic.Anthropic(api_key=claude_api_key)
            
            # 메시지 생성
            message = client.messages.create(
                model="claude-3-haiku-20240307",
                max_tokens=max_tokens,
                temperature=temperature,
                system=system_message,
                messages=[
                    {
                        "role": "user",
                        "content": question
                    }
                ]
            )
            
            return message.content[0].text
                
        except anthropic.APIError as ae:
            return f"Claude API 오류: {str(ae)}"
        except anthropic.RateLimitError:
            return "요청 한도를 초과했습니다. 잠시 후 다시 시도해주세요."
        except Exception as e:
            return f"예상치 못한 오류가 발생했습니다: {str(e)}"

    #################
    # Claude 탭
    #################
    with gr.Tab("claude-3-haiku"):
        gr.Markdown("claude-3-haiku모델")

        # 입력1~5 (한 줄)
        with gr.Row():
            claude_input1 = gr.Textbox(label="입력1", lines=1)
            claude_input2 = gr.Textbox(label="입력2", lines=1)
            claude_input3 = gr.Textbox(label="입력3", lines=1)
            claude_input4 = gr.Textbox(label="입력4", lines=1)
            claude_input5 = gr.Textbox(label="입력5", lines=1)

        # 답변
        claude_answer_output = gr.Textbox(label="답변", interactive=False, lines=5)

        with gr.Accordion("고급 설정 (Claude)", open=False):
            claude_system_message = gr.Textbox(
                label="System Message",
                value="""반드시 한글로 답변할 것.
너는 Anthropic에서 개발한 클로드이다.
최대한 정확하고 친절하게 답변하라.""",
                lines=3
            )
            claude_max_tokens = gr.Slider(
                minimum=1,
                maximum=4096,
                value=512,
                step=1,
                label="Max Tokens"
            )
            claude_temperature = gr.Slider(
                minimum=0.1,
                maximum=2.0,
                value=0.7,
                step=0.05,
                label="Temperature"
            )
            claude_top_p = gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p"
            )

        claude_submit_button = gr.Button("전송")

        def merge_and_call_claude(i1, i2, i3, i4, i5, sys_msg, mt, temp, top_p_, claude_key):
            question = " ".join([i1, i2, i3, i4, i5])
            return respond_claude_qna(
                question=question,
                system_message=sys_msg,
                max_tokens=mt,
                temperature=temp,
                top_p=top_p_,
                claude_api_key=claude_key
            )

        claude_submit_button.click(
            fn=merge_and_call_claude,
            inputs=[
                claude_input1, claude_input2, claude_input3, claude_input4, claude_input5,
                claude_system_message,
                claude_max_tokens,
                claude_temperature,
                claude_top_p,
                claude_token_box
            ],
            outputs=claude_answer_output
        )

#############################
# 메인 실행부
#############################
if __name__ == "__main__":
    demo.launch()