app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
import pandas as pd
from typing import List, Dict, Tuple
import json
import io
import traceback
import csv

# 추론 API 클라이언트 설정
hf_client = InferenceClient(
    "CohereForAI/c4ai-command-r-plus-08-2024", token=os.getenv("HF_TOKEN")
)

def load_code(filename: str) -> str:
    try:
        with open(filename, 'r', encoding='utf-8') as file:
            return file.read()
    except FileNotFoundError:
        return f"{filename} 파일을 찾을 수 없습니다."
    except Exception as e:
        return f"파일을 읽는 중 오류가 발생했습니다: {str(e)}"

def load_parquet(filename: str) -> str:
    try:
        df = pd.read_parquet(filename, engine='pyarrow')
        return df.head(10).to_markdown(index=False)
    except FileNotFoundError:
        return f"{filename} 파일을 찾을 수 없습니다."
    except Exception as e:
        return f"파일을 읽는 중 오류가 발생했습니다: {str(e)}"

def respond(
    message: str,
    history: List[Dict[str, str]],
    system_message: str = "",
    max_tokens: int = 4000,
    temperature: float = 0.5,
    top_p: float = 0.9,
    parquet_data: str = None
) -> str:
    # 시스템 프롬프트 설정
    if parquet_data:
        system_prefix = """반드시 한글로 답변할 것. 너는 업로드된 데이터를 기반으로 질문에 답변하는 역할을 한다. 데이터를 분석하여 사용자에게 도움이 되는 정보를 제공하라. 데이터를 활용하여 상세하고 정확한 답변을 제공하되, 민감한 정보나 개인 정보를 노출하지 마라."""
        try:
            df = pd.read_json(io.StringIO(parquet_data))
            # 데이터의 요약 정보 생성
            data_summary = df.describe(include='all').to_string()
            system_prefix += f"\n\n업로드된 데이터의 요약 정보:\n{data_summary}"
        except Exception as e:
            print(f"데이터 로드 중 오류 발생: {str(e)}\n{traceback.format_exc()}")
            system_prefix += "\n\n데이터를 로드하는 중 오류가 발생했습니다."
    else:
        system_prefix = system_message or "너는 AI 조언자 역할이다."

    # 메시지 생성
    prompt = system_prefix + "\n\n"
    for chat in history:
        if chat['role'] == 'user':
            prompt += f"사용자: {chat['content']}\n"
        else:
            prompt += f"AI: {chat['content']}\n"
    prompt += f"사용자: {message}\nAI:"

    try:
        # 모델에 메시지 전송 및 응답 받기
        response = ""
        stream = hf_client.text_generation(
            prompt=prompt,
            max_new_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        )
        for msg in stream:
            if msg:
                response += msg
                yield response
    except Exception as e:
        error_message = f"추론 중 오류가 발생했습니다: {str(e)}\n{traceback.format_exc()}"
        print(error_message)
        yield error_message

def upload_csv(file_path: str) -> Tuple[str, str]:
    try:
        # CSV 파일 읽기
        df = pd.read_csv(file_path, sep=',')
        # 필수 컬럼 확인
        required_columns = {'id', 'text', 'label', 'metadata'}
        available_columns = set(df.columns)
        missing_columns = required_columns - available_columns
        if missing_columns:
            return f"CSV 파일에 다음 필수 컬럼이 누락되었습니다: {', '.join(missing_columns)}", ""
        # 데이터 클렌징
        df.drop_duplicates(inplace=True)
        df.fillna('', inplace=True)
        # 데이터 유형 최적화
        df = df.astype({'id': 'int32', 'text': 'string', 'label': 'category', 'metadata': 'string'})
        # Parquet 파일로 변환
        parquet_filename = os.path.splitext(os.path.basename(file_path))[0] + '.parquet'
        df.to_parquet(parquet_filename, engine='pyarrow', compression='snappy')
        return f"{parquet_filename} 파일이 성공적으로 업로드되고 변환되었습니다.", parquet_filename
    except Exception as e:
        return f"CSV 파일 업로드 및 변환 중 오류가 발생했습니다: {str(e)}", ""

def upload_parquet(file_path: str) -> Tuple[str, str, str]:
    try:
        # Parquet 파일 읽기
        df = pd.read_parquet(file_path, engine='pyarrow')
        # Markdown으로 변환하여 미리보기
        parquet_content = df.head(10).to_markdown(index=False)
        # DataFrame을 JSON 문자열로 변환
        parquet_json = df.to_json(orient='records', force_ascii=False)
        return "Parquet 파일이 성공적으로 업로드되었습니다.", parquet_content, parquet_json
    except Exception as e:
        return f"Parquet 파일 업로드 중 오류가 발생했습니다: {str(e)}", "", ""

def text_to_parquet(text: str) -> Tuple[str, str, str]:
    try:
        from io import StringIO
        # CSV 데이터를 StringIO를 통해 읽기
        csv_data = StringIO(text)
        df = pd.read_csv(
            csv_data,
            sep=',',
            dtype=str,
            quoting=csv.QUOTE_ALL,  # 모든 필드를 큰따옴표로 감싸는 것으로 처리
            escapechar='\\',        # 이스케이프 문자 설정
            engine='python',        # Python 엔진 사용
            header=None,            # 첫 번째 행을 열 이름으로 사용하지 않음
            names=['id', 'text', 'label', 'metadata']  # 열 이름 지정
        )
        # 데이터 유형 최적화
        df = df.astype({'id': 'int32', 'text': 'string', 'label': 'string', 'metadata': 'string'})
        # Parquet 파일로 변환
        parquet_filename = 'text_to_parquet.parquet'
        df.to_parquet(parquet_filename, engine='pyarrow', compression='snappy')
        # Parquet 파일 내용 미리보기
        parquet_content = load_parquet(parquet_filename)
        return f"{parquet_filename} 파일이 성공적으로 변환되었습니다.", parquet_content, parquet_filename
    except Exception as e:
        error_message = f"텍스트 변환 중 오류가 발생했습니다: {str(e)}\n{traceback.format_exc()}"
        print(error_message)
        return error_message, "", ""


def preprocess_text_with_llm(input_text: str) -> str:
    if not input_text.strip():
        return "입력 텍스트가 비어있습니다."
        
    system_prompt = """당신은 데이터 전처리 전문가입니다. 입력된 텍스트를 CSV 데이터셋 형식으로 변환하세요.

규칙:
1. 출력 형식: id,text,label,metadata
2. id: 1부터 시작하는 순차적 번호
3. text: 의미 있는 단위로 분리된 텍스트
4. label: 텍스트의 주제나 카테고리
5. metadata: 추가 정보(날짜, 출처 등)

주의사항:
- 텍스트에 쉼표가 있으면 큰따옴표로 감싸기
- 큰따옴표는 백슬래시로 이스케이프 처리
- 각 행은 새로운 줄로 구분
- 모든 필드는 쉼표로 구분"""

    full_prompt = f"{system_prompt}\n\n{input_text}\n\n출력:"

    try:
        response = ""
        stream = hf_client.text_generation(
            prompt=full_prompt,
            max_new_tokens=4000,
            temperature=0.3,
            top_p=0.9,
            stream=True,
        )
        
        for msg in stream:
            if msg:
                response += msg
        
        # 응답 정제
        processed_text = response.strip()
        
        # CSV 형식 검증
        try:
            from io import StringIO
            import csv
            csv.reader(StringIO(processed_text))
            return processed_text
        except csv.Error:
            return "LLM이 올바른 CSV 형식을 생성하지 못했습니다. 다시 시도해주세요."
            
    except Exception as e:
        error_message = f"전처리 중 오류가 발생했습니다: {str(e)}"
        print(error_message)
        return error_message


# CSS 설정
css = """
footer {
    visibility: hidden;
}
#chatbot-container, #chatbot-data-upload {
    height: 700px;
    overflow-y: scroll;
}
#chatbot-container .message, #chatbot-data-upload .message {
    font-size: 14px;
}
/* 입력창 배경색 및 글자색 변경 */
textarea, input[type="text"] {
    background-color: #ffffff; /* 흰색 배경 */
    color: #000000; /* 검정색 글자 */
}
/* 파일 업로드 영역 높이 조절 */
#parquet-upload-area {
    max-height: 150px;
    overflow-y: auto;
}
/* 초기 설명 글씨 크기 조절 */
#initial-description {
    font-size: 14px;
}
"""

# Gradio Blocks 인터페이스 설정
with gr.Blocks(css=css) as demo:
    gr.Markdown("# My RAG: LLM이 나만의 데이터로 학습한 콘텐츠 생성/답변", elem_id="initial-description")
    gr.Markdown(
        "### 1) 나만의 데이터를 입력 또는 CSV 업로드로 Parquet 데이터셋 자동 변환 2) Parquet 데이터셋을 업로드하면, LLM이 맞춤 학습 데이터로 활용하여 응답\n"
        "### Tip) '예제'를 통해 다양한 활용 방법을 체험하고 응용해 보세요, 데이터셋 업로드시 미리보기는 10건만 출력",
        elem_id="initial-description"
    )



    # 첫 번째 탭: 챗봇 데이터 업로드 (탭 이름 변경: "My 데이터셋+LLM")
    with gr.Tab("My 데이터셋+LLM"):
        gr.Markdown("### LLM과 대화하기")
        chatbot_data_upload = gr.Chatbot(label="챗봇", type="messages", elem_id="chatbot-data-upload")
        msg_data_upload = gr.Textbox(label="메시지 입력", placeholder="여기에 메시지를 입력하세요...")
        send_data_upload = gr.Button("전송")

        with gr.Accordion("시스템 프롬프트 및 옵션 설정", open=False):
            system_message = gr.Textbox(label="System Message", value="너는 AI 조언자 역할이다.")
            max_tokens = gr.Slider(minimum=1, maximum=8000, value=1000, label="Max Tokens")
            temperature = gr.Slider(minimum=0, maximum=1, value=0.7, label="Temperature")
            top_p = gr.Slider(minimum=0, maximum=1, value=0.9, label="Top P")

        parquet_data_state = gr.State()

        def handle_message_data_upload(
            message: str,
            history: List[Dict[str, str]],
            system_message: str,
            max_tokens: int,
            temperature: float,
            top_p: float,
            parquet_data: str
        ):
            history = history or []
            try:
                # 사용자의 메시지를 히스토리에 추가
                history.append({"role": "user", "content": message})
                # 응답 생성
                response_gen = respond(
                    message, history, system_message, max_tokens, temperature, top_p, parquet_data
                )
                partial_response = ""
                for partial in response_gen:
                    partial_response = partial
                    # 대화 내역 업데이트
                    display_history = history + [
                        {"role": "assistant", "content": partial_response}
                    ]
                    yield display_history, ""
                # 어시스턴트의 응답을 히스토리에 추가
                history.append({"role": "assistant", "content": partial_response})
            except Exception as e:
                response = f"추론 중 오류가 발생했습니다: {str(e)}"
                history.append({"role": "assistant", "content": response})
                yield history, ""

        send_data_upload.click(
            handle_message_data_upload,
            inputs=[
                msg_data_upload,
                chatbot_data_upload,
                system_message,
                max_tokens,
                temperature,
                top_p,
                parquet_data_state,  # parquet_data_state를 사용하여 업로드된 데이터를 전달
            ],
            outputs=[chatbot_data_upload, msg_data_upload],
            queue=True
        )

        # 예제 추가
        with gr.Accordion("예제", open=False):
            gr.Examples(
                examples=[
                    ["업로드된 데이터셋에 대해 요약 설명하라."],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 본 서비스를 SEO 최적화하여 블로그 포스트(개요, 배경 및 필요성, 기존 유사 제품/서비스와 비교하여 특장점, 활용처, 가치, 기대효과, 결론을 포함)로 4000 토큰 이상 작성하라"],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 사용 방법과 차별점, 특징, 강점을 중심으로 4000 토큰 이상 유튜브 영상 스크립트 형태로 작성하라"],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 제품 상세 페이지 형식의 내용을 4000 토큰 이상 자세히 설명하라"],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, FAQ 20건을 상세하게 작성하라. 4000토큰 이상 사용하라."],
                    ["업로드된 데이터셋 파일을 학습 데이터로 활용하여, 특허 출원에 활용할 기술 및 비즈니스 모델 측면을 포함하여 특허 출원서 구성에 맞게 혁신적인 창의 발명 내용을 중심으로 4000 토큰 이상 작성하라."],
                ],
                inputs=msg_data_upload,
                label="예제 선택",
            )

        # Parquet 파일 업로드를 화면 하단으로 이동
        gr.Markdown("### Parquet 파일 업로드")
        with gr.Row():
            with gr.Column():
                parquet_upload = gr.File(
                    label="Parquet 파일 업로드", type="filepath", elem_id="parquet-upload-area"
                )
                parquet_upload_button = gr.Button("업로드")
                parquet_upload_status = gr.Textbox(label="업로드 상태", interactive=False)
                parquet_preview_chat = gr.Markdown(label="Parquet 파일 미리보기")

                def handle_parquet_upload(file_path: str):
                    message, parquet_content, parquet_json = upload_parquet(file_path)
                    if parquet_json:
                        return message, parquet_content, parquet_json
                    else:
                        return message, "", ""

                parquet_upload_button.click(
                    handle_parquet_upload,
                    inputs=parquet_upload,
                    outputs=[parquet_upload_status, parquet_preview_chat, parquet_data_state]
                )

    # 두 번째 탭: 데이터 변환 (탭 이름 변경: "CSV to My 데이터셋")
    with gr.Tab("CSV to My 데이터셋"):
        gr.Markdown("### CSV 파일 업로드 및 Parquet 변환")
        with gr.Row():
            with gr.Column():
                csv_file = gr.File(label="CSV 파일 업로드", type="filepath")
                upload_button = gr.Button("업로드 및 변환")
                upload_status = gr.Textbox(label="업로드 상태", interactive=False)
                parquet_preview = gr.Markdown(label="Parquet 파일 미리보기")
                download_button = gr.File(label="Parquet 파일 다운로드", interactive=False)

                def handle_csv_upload(file_path: str):
                    message, parquet_filename = upload_csv(file_path)
                    if parquet_filename:
                        parquet_content = load_parquet(parquet_filename)
                        return message, parquet_content, parquet_filename
                    else:
                        return message, "", None

                upload_button.click(
                    handle_csv_upload,
                    inputs=csv_file,
                    outputs=[upload_status, parquet_preview, download_button]
                )

    # 세 번째 탭: 텍스트 to csv to parquet 변환 (탭 이름 변경: "Text to My 데이터셋")
    with gr.Tab("Text to My 데이터셋"):
        gr.Markdown("### 텍스트를 입력하면 CSV로 변환 후 Parquet으로 자동 전환됩니다.")
        with gr.Row():
            with gr.Column():
                text_input = gr.Textbox(
                    label="텍스트 입력 (각 행은 `id,text,label,metadata` 형식으로 입력)",
                    lines=10,
                    placeholder='예: 1,"이순신","장군","거북선"\n2,"원균","장군","모함"\n3,"선조","왕","시기"\n4,"도요토미 히데요시","왕","침략"'
                )
                convert_button = gr.Button("변환 및 다운로드")
                convert_status = gr.Textbox(label="변환 상태", interactive=False)
                parquet_preview_convert = gr.Markdown(label="Parquet 파일 미리보기")
                download_parquet_convert = gr.File(label="Parquet 파일 다운로드", interactive=False)

                def handle_text_to_parquet(text: str):
                    message, parquet_content, parquet_filename = text_to_parquet(text)
                    if parquet_filename:
                        return message, parquet_content, parquet_filename
                    else:
                        return message, "", None

                convert_button.click(
                    handle_text_to_parquet,
                    inputs=text_input,
                    outputs=[convert_status, parquet_preview_convert, download_parquet_convert]
                )

    # 네 번째 탭: 텍스트를 데이터셋 형식으로 전처리 (개선된 버전)
    with gr.Tab("Text Preprocessing with LLM"):
        gr.Markdown("### 텍스트를 입력하면 LLM이 데이터셋 형식에 맞게 전처리하여 출력합니다.")
        with gr.Row():
            with gr.Column():
                raw_text_input = gr.Textbox(
                    label="텍스트 입력",
                    lines=15,
                    placeholder="여기에 전처리할 텍스트를 입력하세요..."
                )
                
                with gr.Row():
                    preprocess_button = gr.Button("전처리 실행", variant="primary")
                    clear_button = gr.Button("초기화")

                preprocess_status = gr.Textbox(
                    label="전처리 상태",
                    interactive=False,
                    value="대기 중..."
                )
                
                processed_text_output = gr.Textbox(
                    label="전처리된 데이터셋 출력",
                    lines=15,
                    interactive=False
                )
                
                # Parquet 변환 및 다운로드 섹션
                with gr.Row():
                    convert_to_parquet_button = gr.Button("Parquet으로 변환", visible=True)
                    download_parquet = gr.File(
                        label="변환된 Parquet 파일 다운로드",
                        visible=False
                    )

                def handle_text_preprocessing(input_text: str):
                    if not input_text.strip():
                        return "입력 텍스트가 없습니다.", ""
                    
                    try:
                        preprocess_status_msg = "전처리를 시작합니다..."
                        yield preprocess_status_msg, ""
                        
                        processed_text = preprocess_text_with_llm(input_text)
                        
                        if processed_text:
                            preprocess_status_msg = "전처리가 완료되었습니다."
                            yield preprocess_status_msg, processed_text
                        else:
                            preprocess_status_msg = "전처리 결과가 없습니다."
                            yield preprocess_status_msg, ""
                            
                    except Exception as e:
                        error_msg = f"처리 중 오류가 발생했습니다: {str(e)}"
                        yield error_msg, ""

                def clear_inputs():
                    return "", "대기 중...", ""

                def convert_to_parquet_file(processed_text: str):
                    if not processed_text.strip():
                        return "변환할 텍스트가 없습니다.", None
                    
                    try:
                        message, parquet_content, parquet_filename = text_to_parquet(processed_text)
                        if parquet_filename:
                            return message, parquet_filename
                        return message, None
                    except Exception as e:
                        return f"Parquet 변환 중 오류 발생: {str(e)}", None

                # 이벤트 핸들러 연결
                preprocess_button.click(
                    handle_text_preprocessing,
                    inputs=[raw_text_input],
                    outputs=[preprocess_status, processed_text_output],
                    queue=True
                )

                clear_button.click(
                    clear_inputs,
                    outputs=[raw_text_input, preprocess_status, processed_text_output]
                )

                convert_to_parquet_button.click(
                    convert_to_parquet_file,
                    inputs=[processed_text_output],
                    outputs=[preprocess_status, download_parquet]
                )

                # 예제 텍스트 추가
                with gr.Accordion("예제 텍스트", open=False):
                    gr.Examples(
                        examples=[
                            ["이순신은 조선 중기의 무신이다. 그는 임진왜란 당시 해군을 이끌었다. 거북선을 만들어 왜군과 싸웠다."],
                            ["인공지능은 컴퓨터 과학의 한 분야이다. 기계학습은 인공지능의 하위 분야이다. 딥러닝은 기계학습의 한 방법이다."]
                        ],
                        inputs=raw_text_input,
                        label="예제 선택"
                    )

    gr.Markdown("### [email protected]", elem_id="initial-description")

if __name__ == "__main__":
    demo.launch(share=True)