import gradio as gr
import random
import os
import pandas as pd  # 엑셀 파일 처리를 위한 pandas 라이브러리 추가
from huggingface_hub import InferenceClient

MODELS = {
    "Zephyr 7B Beta": "HuggingFaceH4/zephyr-7b-beta",
    "DeepSeek Coder V2": "deepseek-ai/DeepSeek-Coder-V2-Instruct",
    "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",
    "Cohere Command R+": "CohereForAI/c4ai-command-r-plus",
    "Aya-23-35B": "CohereForAI/aya-23-35B"
}

def create_client(model_name):
    return InferenceClient(model_name, token=os.getenv("HF_TOKEN"))

def call_api(model, content, system_message, max_tokens, temperature, top_p):
    client = create_client(MODELS[model])
    messages = [{"role": "system", "content": system_message}, {"role": "user", "content": content}]
    random_seed = random.randint(0, 1000000)
    response = client.chat_completion(messages=messages, max_tokens=max_tokens, temperature=temperature, top_p=top_p, seed=random_seed)
    return response.choices[0].message.content

def generate_text(model, user_message, system_message, max_tokens, temperature, top_p):
    return call_api(model, user_message, system_message, max_tokens, temperature, top_p)

# 엑셀 파일을 처리하는 함수 추가
def process_excel(file):
    if file is not None:
        df = pd.read_excel(file.name)  # 엑셀 파일을 읽음
        return df.head().to_string()  # 엑셀 파일의 처음 몇 줄을 출력
    else:
        return "엑셀 파일이 업로드되지 않았습니다."

title = "AI 텍스트 생성기"

with gr.Blocks() as demo:
    gr.Markdown(f"# {title}")
    
    model = gr.Radio(choices=list(MODELS.keys()), label="언어 모델 선택", value="Zephyr 7B Beta")
    user_message = gr.Textbox(label="사용자 메시지", lines=5)
    system_message = gr.Textbox(label="시스템 메시지 (프롬프트)", lines=10)
    
    with gr.Accordion("고급 설정", open=False):
        max_tokens = gr.Slider(label="Max Tokens", minimum=0, maximum=4000, value=500, step=100)
        temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.75, step=0.05)
        top_p = gr.Slider(label="Top P", minimum=0.1, maximum=1.0, value=0.95, step=0.05)
    
    # 엑셀 파일 업로드 컴포넌트 추가
    excel_file = gr.File(label="엑셀 파일 업로드", type="file")
    excel_output = gr.Textbox(label="엑셀 내용", lines=10)

    generate_btn = gr.Button("텍스트 생성하기")
    output = gr.Textbox(label="생성된 텍스트", lines=10)
    
    generate_btn.click(fn=generate_text, 
                       inputs=[model, user_message, system_message, max_tokens, temperature, top_p], 
                       outputs=[output])
    
    # 엑셀 파일 업로드 버튼과 처리 함수 연결
    excel_file.upload(fn=process_excel, inputs=[excel_file], outputs=[excel_output])

demo.launch()