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:
    # 시스템 프롬프트에 중복 방지 지시 추가
    system_prefix = """반드시 한글로 답변할 것. 너는 업로드된 데이터를 기반으로 질문에 답변하는 역할을 한다.
    
    중요 규칙:
    1. 이전 대화에서 이미 답변한 내용을 반복하지 말 것
    2. 질문과 직접 관련된 내용만 답변할 것
    3. 불필요한 예시나 부연 설명은 최소화할 것
    4. 답변은 명확하고 간결하게 할 것
    5. 동일한 내용을 다른 표현으로 반복하지 말 것
    """
    
    if parquet_data:
        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)}")
    
    # 이전 대화 컨텍스트 최적화
    recent_history = history[-3:] if history else []  # 최근 3개 대화만 유지
    
    prompt = system_prefix + "\n\n"
    for chat in recent_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,
            repetition_penalty=1.2,  # 반복 페널티 추가
            no_repeat_ngram_size=3,  # n-gram 반복 방지
        )
        
        for msg in stream:
            if msg:
                response += msg
                # 중복 문장 제거
                response = remove_duplicates(response)
                yield response
    except Exception as e:
        error_message = f"추론 오류: {str(e)}"
        print(error_message)
        yield error_message

def remove_duplicates(text: str) -> str:
    """중복 문장 제거 함수"""
    sentences = text.split('.')
    unique_sentences = []
    seen = set()
    
    for sentence in sentences:
        sentence = sentence.strip()
        if sentence and sentence not in seen:
            seen.add(sentence)
            unique_sentences.append(sentence)
    
    return '. '.join(unique_sentences)

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
        import csv
        
        # 입력 텍스트 정제
        lines = text.strip().split('\n')
        cleaned_lines = []
        
        for line in lines:
            # 빈 줄 건너뛰기
            if not line.strip():
                continue
            
            # 쌍따옴표 정규화
            line = line.replace('""', '"')  # 중복 쌍따옴표 처리
            
            # CSV 파싱을 위한 임시 StringIO 객체 생성
            temp_buffer = StringIO(line)
            try:
                # CSV 라인 파싱 시도
                reader = csv.reader(temp_buffer, quoting=csv.QUOTE_ALL)
                parsed_line = next(reader)
                if len(parsed_line) == 4:  # id, text, label, metadata
                    # 각 필드를 적절히 포맷팅
                    formatted_line = f'{parsed_line[0]},"{parsed_line[1]}","{parsed_line[2]}","{parsed_line[3]}"'
                    cleaned_lines.append(formatted_line)
            except:
                continue
            finally:
                temp_buffer.close()
        
        # 정제된 CSV 데이터 생성
        cleaned_csv = '\n'.join(cleaned_lines)
        
        # DataFrame 생성
        df = pd.read_csv(
            StringIO(cleaned_csv),
            sep=',',
            quoting=csv.QUOTE_ALL,
            escapechar='\\',
            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)}"
        print(f"{error_message}\n{traceback.format_exc()}")
        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: 텍스트의 주제나 카테고리를 아래 기준으로 정확하게 한 개만 선택
   - Historical_Figure (역사적 인물)
   - Military_History (군사 역사)
   - Technology (기술)
   - Politics (정치)
   - Culture (문화)
5. metadata: 날짜, 출처 등 추가 정보

중요:
- 동일한 텍스트를 반복해서 출력하지 말 것
- 각 텍스트는 한 번만 처리하여 가장 적합한 label을 선택할 것
- 입력 텍스트를 의미 단위로 적절히 분리할 것

예시:
1,"이순신은 조선 중기의 무신이다.","Historical_Figure","조선시대, 위키백과"

주의사항:
- text에 쉼표가 있으면 큰따옴표로 감싸기
- 큰따옴표는 백슬래시로 이스케이프 처리
- 각 행은 새로운 줄로 구분
- 불필요한 반복 출력 금지"""

    full_prompt = f"{system_prompt}\n\n입력텍스트:\n{input_text}\n\n출력:"

    try:
        response = ""
        stream = hf_client.text_generation(
            prompt=full_prompt,
            max_new_tokens=4000,
            temperature=0.1,  # 더 결정적인 출력을 위해 낮춤
            top_p=0.9,
            stream=True,
        )
        
        for msg in stream:
            if msg:
                response += msg
        
        # <EOS_TOKEN> 이전까지만 추출하고 정제
        if "<EOS_TOKEN>" in response:
            processed_text = response.split("<EOS_TOKEN>")[0].strip()
        else:
            processed_text = response.strip()
            
        # 중복 출력 제거
        lines = processed_text.split('\n')
        unique_lines = []
        seen_texts = set()
        
        for line in lines:
            line = line.strip()
            if line and '출력:' not in line and line not in seen_texts:
                unique_lines.append(line)
                seen_texts.add(line)
        
        processed_text = '\n'.join(unique_lines)
        
        # 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 []
    
            # 중복 질문 체크
            if history and any(chat['role'] == 'user' and chat['content'].strip() == message.strip() for chat in history[-3:]):
                yield history + [{"role": "assistant", "content": "동일한 질문이 최근에 있었습니다. 다른 질문을 해주세요."}], ""
                return
    
            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": remove_duplicates(partial_response)}]
                    yield display_history, ""
            
                history.append({"role": "assistant", "content": remove_duplicates(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]
                )

    # 네번째 탭의 UI 부분 수정
    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 변환 및 다운로드 섹션
                convert_to_parquet_button = gr.Button("Parquet으로 변환")
                download_parquet = gr.File(label="변환된 Parquet 파일 다운로드")




                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)