Create app.py

app.py

@@ -0,0 +1,343 @@
+import torch
+import gradio as gr
+import spaces
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer
+import os, gc, logging
+from threading import Thread
+import random
+from datasets import load_dataset
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+import pandas as pd
+from typing import List, Tuple, Iterator
+import json
+from datetime import datetime
+from concurrent.futures import ThreadPoolExecutor
+from functools import lru_cache
+import pyarrow.parquet as pq
+import pypdf
+from pdfminer.high_level import extract_text
+from pdfminer.layout import LAParams
+from tabulate import tabulate
+from pydantic import BaseModel
+import unittest
+
+# 로깅 설정
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.FileHandler('app.log'),
+        logging.StreamHandler()
+    ]
+)
+logger = logging.getLogger(__name__)
+
+# 설정 클래스
+class Config:
+    def __init__(self):
+        self.MODEL_ID = "CohereForAI/c4ai-command-r7b-12-2024"
+        self.MAX_HISTORY = 10
+        self.MAX_TOKENS = 4096
+        self.DEFAULT_TEMPERATURE = 0.8
+        self.HF_TOKEN = os.environ.get("HF_TOKEN", None)
+        self.MODELS = os.environ.get("MODELS")
+        
+config = Config()
+
+# 커스텀 예외 클래스
+class FileProcessingError(Exception):
+    pass
+
+# 응답 모델
+class ChatResponse(BaseModel):
+    message: str
+    status: str
+    timestamp: datetime
+
+# 파일 처리 클래스
+class FileProcessor:
+    @staticmethod
+    def process_pdf(file_path):
+        try:
+            with ThreadPoolExecutor() as executor:
+                pdf_reader = pypdf.PdfReader(file_path)
+                text = extract_text(
+                    file_path,
+                    laparams=LAParams(
+                        line_margin=0.5,
+                        word_margin=0.1,
+                        char_margin=2.0,
+                        all_texts=True
+                    )
+                )
+                return text
+        except Exception as e:
+            raise FileProcessingError(f"PDF processing error: {str(e)}")
+
+    @staticmethod
+    def process_csv(file_path):
+        try:
+            encodings = ['utf-8', 'cp949', 'euc-kr', 'latin1']
+            for encoding in encodings:
+                try:
+                    return pd.read_csv(file_path, encoding=encoding)
+                except UnicodeDecodeError:
+                    continue
+            raise FileProcessingError("Unable to read CSV with supported encodings")
+        except Exception as e:
+            raise FileProcessingError(f"CSV processing error: {str(e)}")
+
+# 메모리 관리
+@torch.no_grad()
+def clear_cuda_memory():
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        gc.collect()
+
+# 모델 로드
+@spaces.GPU
+def load_model():
+    try:
+        model = AutoModelForCausalLM.from_pretrained(
+            config.MODEL_ID,
+            torch_dtype=torch.bfloat16,
+            device_map="auto",
+        )
+        return model
+    except Exception as e:
+        logger.error(f"Model loading error: {str(e)}")
+        raise
+
+# 컨텍스트 검색
+@lru_cache(maxsize=100)
+def find_relevant_context(query, top_k=3):
+    try:
+        query_vector = vectorizer.transform([query])
+        similarities = (query_vector * question_vectors.T).toarray()[0]
+        top_indices = np.argsort(similarities)[-top_k:][::-1]
+        
+        relevant_contexts = []
+        for idx in top_indices:
+            if similarities[idx] > 0:
+                relevant_contexts.append({
+                    'question': questions[idx],
+                    'answer': wiki_dataset['train']['answer'][idx],
+                    'similarity': similarities[idx]
+                })
+        return relevant_contexts
+    except Exception as e:
+        logger.error(f"Context search error: {str(e)}")
+        return []
+
+# 스트리밍 채팅
+@spaces.GPU
+def stream_chat(message: str, history: list, uploaded_file, temperature: float, 
+                max_new_tokens: int, top_p: float, top_k: int, penalty: float) -> Iterator[Tuple[str, list]]:
+    """
+    스트리밍 채팅 응답을 생성합니다.
+    
+    Args:
+        message (str): 사용자 입력 메시지
+        history (list): 대화 히스토리
+        uploaded_file: 업로드된 파일
+        temperature (float): 생성 온도
+        max_new_tokens (int): 최대 토큰 수
+        top_p (float): 상위 p 샘플링
+        top_k (int): 상위 k 샘플링
+        penalty (float): 반복 페널티
+        
+    Returns:
+        Iterator[Tuple[str, list]]: 생성된 응답과 업데이트된 히스토리
+    """
+    global model, current_file_context
+    
+    try:
+        if model is None:
+            model = load_model()
+            
+        logger.info(f'Processing message: {message}')
+        logger.debug(f'History length: {len(history)}')
+
+        # 파일 처리
+        file_context = ""
+        if uploaded_file:
+            try:
+                file_ext = os.path.splitext(uploaded_file.name)[1].lower()
+                if file_ext == '.pdf':
+                    content = FileProcessor.process_pdf(uploaded_file.name)
+                elif file_ext == '.csv':
+                    content = FileProcessor.process_csv(uploaded_file.name)
+                else:
+                    content = safe_file_read(uploaded_file.name)
+                
+                file_context = analyze_file_content(content, file_ext)
+                current_file_context = file_context
+            except Exception as e:
+                logger.error(f"File processing error: {str(e)}")
+                file_context = f"\n\n❌ File analysis error: {str(e)}"
+
+        # 컨텍스트 검색 및 프롬프트 구성
+        relevant_contexts = find_relevant_context(message)
+        wiki_context = "\n\n관련 위키피디아 정보:\n" + "\n".join([
+            f"Q: {ctx['question']}\nA: {ctx['answer']}\n유사도: {ctx['similarity']:.3f}"
+            for ctx in relevant_contexts
+        ])
+
+        # 토큰화 및 생성
+        conversation = [
+            {"role": "user" if i % 2 == 0 else "assistant", "content": msg}
+            for hist in history[-config.MAX_HISTORY:]
+            for i, msg in enumerate(hist)
+        ]
+        
+        final_message = f"{file_context}{wiki_context}\n현재 질문: {message}"
+        conversation.append({"role": "user", "content": final_message})
+
+        inputs = tokenizer(
+            tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True),
+            return_tensors="pt"
+        ).to("cuda")
+
+        streamer = TextIteratorStreamer(tokenizer, timeout=10., skip_prompt=True, skip_special_tokens=True)
+        
+        generate_kwargs = dict(
+            inputs, 
+            streamer=streamer,
+            top_k=top_k,
+            top_p=top_p,
+            repetition_penalty=penalty,
+            max_new_tokens=min(max_new_tokens, 2048),
+            do_sample=True, 
+            temperature=temperature,
+            eos_token_id=[255001],
+        )
+        
+        clear_cuda_memory()
+        
+        thread = Thread(target=model.generate, kwargs=generate_kwargs)
+        thread.start()
+
+        buffer = ""
+        for new_text in streamer:
+            buffer += new_text
+            yield "", history + [[message, buffer]]
+
+        clear_cuda_memory()
+
+    except Exception as e:
+        logger.error(f"Stream chat error: {str(e)}")
+        yield "", history + [[message, f"Error: {str(e)}"]]
+        clear_cuda_memory()
+
+# UI 생성
+def create_demo():
+    with gr.Blocks(css=UPDATED_CSS) as demo:
+        # UI 컴포넌트 구성
+        with gr.Column(elem_classes="markdown-style"):
+            gr.Markdown("""
+                # 🤖 RAGOndevice
+                #### 📊 RAG: Upload and Analyze Files (TXT, CSV, PDF, Parquet files)
+                Upload your files for data analysis and learning
+            """)
+        
+        chatbot = gr.Chatbot(
+            value=[],
+            height=600,
+            label="GiniGEN AI Assistant",
+            elem_classes="chat-container"
+        )
+        
+        # 입력 컴포넌트
+        with gr.Row(elem_classes="input-container"):
+            with gr.Column(scale=1, min_width=70):
+                file_upload = gr.File(
+                    type="filepath",
+                    elem_classes="file-upload-icon",
+                    scale=1,
+                    container=True,
+                    interactive=True,
+                    show_label=False
+                )
+            
+            with gr.Column(scale=3):
+                msg = gr.Textbox(
+                    show_label=False,
+                    placeholder="Type your message here... 💭",
+                    container=False,
+                    elem_classes="input-textbox",
+                    scale=1
+                )
+            
+            with gr.Column(scale=1, min_width=70):
+                send = gr.Button(
+                    "Send",
+                    elem_classes="send-button custom-button",
+                    scale=1
+                )
+                
+            with gr.Column(scale=1, min_width=70):
+                clear = gr.Button(
+                    "Clear",
+                    elem_classes="clear-button custom-button",
+                    scale=1
+                )
+        
+        # 고급 설정
+        with gr.Accordion("🎮 Advanced Settings", open=False):
+            with gr.Row():
+                with gr.Column(scale=1):
+                    temperature = gr.Slider(
+                        minimum=0, maximum=1, step=0.1, value=config.DEFAULT_TEMPERATURE,
+                        label="Creativity Level 🎨"
+                    )
+                    max_new_tokens = gr.Slider(
+                        minimum=128, maximum=8000, step=1, value=4000,
+                        label="Maximum Token Count 📝"
+                    )
+                with gr.Column(scale=1):
+                    top_p = gr.Slider(
+                        minimum=0.0, maximum=1.0, step=0.1, value=0.8,
+                        label="Diversity Control 🎯"
+                    )
+                    top_k = gr.Slider(
+                        minimum=1, maximum=20, step=1, value=20,
+                        label="Selection Range 📊"
+                    )
+                    penalty = gr.Slider(
+                        minimum=0.0, maximum=2.0, step=0.1, value=1.0,
+                        label="Repetition Penalty 🔄"
+                    )
+
+        # 이벤트 바인딩
+        msg.submit(stream_chat, [msg, chatbot, file_upload, temperature, max_new_tokens, top_p, top_k, penalty], [msg, chatbot])
+        send.click(stream_chat, [msg, chatbot, file_upload, temperature, max_new_tokens, top_p, top_k, penalty], [msg, chatbot])
+        clear.click(lambda: ([], None, ""), outputs=[chatbot, file_upload, msg])
+
+        return demo
+
+# 메인 실행
+if __name__ == "__main__":
+    # 위키피디아 데이터셋 로드
+    wiki_dataset = load_dataset("lcw99/wikipedia-korean-20240501-1million-qna")
+    logger.info("Wikipedia dataset loaded")
+
+    # TF-IDF 벡터라이저 초기화
+    questions = wiki_dataset['train']['question'][:10000]
+    vectorizer = TfidfVectorizer(max_features=1000)
+    question_vectors = vectorizer.fit_transform(questions)
+    logger.info("TF-IDF vectorization completed")
+
+    # UI 실행
+    demo = create_demo()
+    demo.launch()
+
+# 테스트 코드
+class TestChatBot(unittest.TestCase):
+    def test_file_processing(self):
+        # 테스트 구현
+        pass
+
+    def test_context_search(self):
+        # 테스트 구현
+        pass